1 //! Rand implementations for complex numbers
2 
3 use crate::Complex;
4 use num_traits::Num;
5 use rand::distributions::Standard;
6 use rand::prelude::*;
7 
8 impl<T> Distribution<Complex<T>> for Standard
9 where
10     T: Num + Clone,
11     Standard: Distribution<T>,
12 {
sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Complex<T>13     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Complex<T> {
14         Complex::new(self.sample(rng), self.sample(rng))
15     }
16 }
17 
18 /// A generic random value distribution for complex numbers.
19 #[derive(Clone, Copy, Debug)]
20 pub struct ComplexDistribution<Re, Im = Re> {
21     re: Re,
22     im: Im,
23 }
24 
25 impl<Re, Im> ComplexDistribution<Re, Im> {
26     /// Creates a complex distribution from independent
27     /// distributions of the real and imaginary parts.
new(re: Re, im: Im) -> Self28     pub fn new(re: Re, im: Im) -> Self {
29         ComplexDistribution { re, im }
30     }
31 }
32 
33 impl<T, Re, Im> Distribution<Complex<T>> for ComplexDistribution<Re, Im>
34 where
35     T: Num + Clone,
36     Re: Distribution<T>,
37     Im: Distribution<T>,
38 {
sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Complex<T>39     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Complex<T> {
40         Complex::new(self.re.sample(rng), self.im.sample(rng))
41     }
42 }
43 
44 #[cfg(test)]
test_rng() -> impl RngCore45 fn test_rng() -> impl RngCore {
46     /// Simple `Rng` for testing without additional dependencies
47     struct XorShiftStar {
48         a: u64,
49     }
50 
51     impl RngCore for XorShiftStar {
52         fn next_u32(&mut self) -> u32 {
53             self.next_u64() as u32
54         }
55 
56         fn next_u64(&mut self) -> u64 {
57             // https://en.wikipedia.org/wiki/Xorshift#xorshift*
58             self.a ^= self.a >> 12;
59             self.a ^= self.a << 25;
60             self.a ^= self.a >> 27;
61             self.a.wrapping_mul(0x2545_F491_4F6C_DD1D)
62         }
63 
64         fn fill_bytes(&mut self, dest: &mut [u8]) {
65             for chunk in dest.chunks_mut(8) {
66                 let bytes = self.next_u64().to_le_bytes();
67                 let slice = &bytes[..chunk.len()];
68                 chunk.copy_from_slice(slice)
69             }
70         }
71 
72         fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
73             self.fill_bytes(dest);
74             Ok(())
75         }
76     }
77 
78     XorShiftStar {
79         a: 0x0123_4567_89AB_CDEF,
80     }
81 }
82 
83 #[test]
standard_f64()84 fn standard_f64() {
85     let mut rng = test_rng();
86     for _ in 0..100 {
87         let c: Complex<f64> = rng.gen();
88         assert!(c.re >= 0.0 && c.re < 1.0);
89         assert!(c.im >= 0.0 && c.im < 1.0);
90     }
91 }
92 
93 #[test]
generic_standard_f64()94 fn generic_standard_f64() {
95     let mut rng = test_rng();
96     let dist = ComplexDistribution::new(Standard, Standard);
97     for _ in 0..100 {
98         let c: Complex<f64> = rng.sample(dist);
99         assert!(c.re >= 0.0 && c.re < 1.0);
100         assert!(c.im >= 0.0 && c.im < 1.0);
101     }
102 }
103 
104 #[test]
generic_uniform_f64()105 fn generic_uniform_f64() {
106     use rand::distributions::Uniform;
107 
108     let mut rng = test_rng();
109     let re = Uniform::new(-100.0, 0.0);
110     let im = Uniform::new(0.0, 100.0);
111     let dist = ComplexDistribution::new(re, im);
112     for _ in 0..100 {
113         // no type annotation required, since `Uniform` only produces one type.
114         let c = rng.sample(dist);
115         assert!(c.re >= -100.0 && c.re < 0.0);
116         assert!(c.im >= 0.0 && c.im < 100.0);
117     }
118 }
119 
120 #[test]
generic_mixed_f64()121 fn generic_mixed_f64() {
122     use rand::distributions::Uniform;
123 
124     let mut rng = test_rng();
125     let re = Uniform::new(-100.0, 0.0);
126     let dist = ComplexDistribution::new(re, Standard);
127     for _ in 0..100 {
128         // no type annotation required, since `Uniform` only produces one type.
129         let c = rng.sample(dist);
130         assert!(c.re >= -100.0 && c.re < 0.0);
131         assert!(c.im >= 0.0 && c.im < 1.0);
132     }
133 }
134 
135 #[test]
generic_uniform_i32()136 fn generic_uniform_i32() {
137     use rand::distributions::Uniform;
138 
139     let mut rng = test_rng();
140     let re = Uniform::new(-100, 0);
141     let im = Uniform::new(0, 100);
142     let dist = ComplexDistribution::new(re, im);
143     for _ in 0..100 {
144         // no type annotation required, since `Uniform` only produces one type.
145         let c = rng.sample(dist);
146         assert!(c.re >= -100 && c.re < 0);
147         assert!(c.im >= 0 && c.im < 100);
148     }
149 }
150