1use 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#[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 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 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 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]
84fn 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]
94fn 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]
105fn 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 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]
121fn 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 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]
136fn 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 let c = rng.sample(dist);
146 assert!(c.re >= -100 && c.re < 0);
147 assert!(c.im >= 0 && c.im < 100);
148 }
149}