num_complex/
crand.rs

1//! Rand implementations for complex numbers
2
3use crate::Complex;
4use num_traits::Num;
5use rand::distributions::Standard;
6use rand::prelude::*;
7
8impl<T> Distribution<Complex<T>> for Standard
9where
10    T: Num + Clone,
11    Standard: Distribution<T>,
12{
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)]
20pub struct ComplexDistribution<Re, Im = Re> {
21    re: Re,
22    im: Im,
23}
24
25impl<Re, Im> ComplexDistribution<Re, Im> {
26    /// Creates a complex distribution from independent
27    /// distributions of the real and imaginary parts.
28    pub fn new(re: Re, im: Im) -> Self {
29        ComplexDistribution { re, im }
30    }
31}
32
33impl<T, Re, Im> Distribution<Complex<T>> for ComplexDistribution<Re, Im>
34where
35    T: Num + Clone,
36    Re: Distribution<T>,
37    Im: Distribution<T>,
38{
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)]
45fn 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            Ok(self.fill_bytes(dest))
74        }
75    }
76
77    XorShiftStar {
78        a: 0x0123_4567_89AB_CDEF,
79    }
80}
81
82#[test]
83fn standard_f64() {
84    let mut rng = test_rng();
85    for _ in 0..100 {
86        let c: Complex<f64> = rng.gen();
87        assert!(c.re >= 0.0 && c.re < 1.0);
88        assert!(c.im >= 0.0 && c.im < 1.0);
89    }
90}
91
92#[test]
93fn generic_standard_f64() {
94    let mut rng = test_rng();
95    let dist = ComplexDistribution::new(Standard, Standard);
96    for _ in 0..100 {
97        let c: Complex<f64> = rng.sample(&dist);
98        assert!(c.re >= 0.0 && c.re < 1.0);
99        assert!(c.im >= 0.0 && c.im < 1.0);
100    }
101}
102
103#[test]
104fn generic_uniform_f64() {
105    use rand::distributions::Uniform;
106
107    let mut rng = test_rng();
108    let re = Uniform::new(-100.0, 0.0);
109    let im = Uniform::new(0.0, 100.0);
110    let dist = ComplexDistribution::new(re, im);
111    for _ in 0..100 {
112        // no type annotation required, since `Uniform` only produces one type.
113        let c = rng.sample(&dist);
114        assert!(c.re >= -100.0 && c.re < 0.0);
115        assert!(c.im >= 0.0 && c.im < 100.0);
116    }
117}
118
119#[test]
120fn generic_mixed_f64() {
121    use rand::distributions::Uniform;
122
123    let mut rng = test_rng();
124    let re = Uniform::new(-100.0, 0.0);
125    let dist = ComplexDistribution::new(re, Standard);
126    for _ in 0..100 {
127        // no type annotation required, since `Uniform` only produces one type.
128        let c = rng.sample(&dist);
129        assert!(c.re >= -100.0 && c.re < 0.0);
130        assert!(c.im >= 0.0 && c.im < 1.0);
131    }
132}
133
134#[test]
135fn generic_uniform_i32() {
136    use rand::distributions::Uniform;
137
138    let mut rng = test_rng();
139    let re = Uniform::new(-100, 0);
140    let im = Uniform::new(0, 100);
141    let dist = ComplexDistribution::new(re, im);
142    for _ in 0..100 {
143        // no type annotation required, since `Uniform` only produces one type.
144        let c = rng.sample(&dist);
145        assert!(c.re >= -100 && c.re < 0);
146        assert!(c.im >= 0 && c.im < 100);
147    }
148}