1//! Complex
  2
  3use rand::Rng;
  4use std::ops::{Add, Div, Mul, Neg, Sub};
  5
  6#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
  7pub struct C<T> {
  8    pub re: T,
  9    pub im: T,
 10}
 11
 12impl From<f64> for C<f64> {
 13    fn from(v: f64) -> Self {
 14        Self { re: v, im: 0_f64 }
 15    }
 16}
 17
 18impl<T> C<T>
 19where
 20    T: Default + From<i32>,
 21{
 22    pub fn new(re: T, im: T) -> Self {
 23        Self { re, im }
 24    }
 25
 26    pub fn rand(mut rng: impl Rng, max: u64) -> Self {
 27        Self::new(
 28            T::from(rng.gen_range(0..max) as i32),
 29            T::from(rng.gen_range(0..max) as i32),
 30        )
 31    }
 32
 33    pub fn zero() -> C<T> {
 34        Self {
 35            re: T::from(0),
 36            im: T::from(0),
 37        }
 38    }
 39    pub fn one() -> C<T> {
 40        Self {
 41            re: T::from(1),
 42            im: T::from(0),
 43        }
 44    }
 45    pub fn i() -> C<T> {
 46        Self {
 47            re: T::from(0),
 48            im: T::from(1),
 49        }
 50    }
 51}
 52
 53impl C<f64> {
 54    // cos & sin from Taylor series approximation, details at
 55    // https://en.wikipedia.org/wiki/Sine_and_cosine#Series_and_polynomials
 56    fn cos(x: f64) -> f64 {
 57        let mut r = 1.0;
 58        let mut term = 1.0;
 59        let mut n = 1;
 60        for _ in 0..10 {
 61            term *= -(x * x) / ((2 * n - 1) * (2 * n)) as f64;
 62            r += term;
 63            n += 1;
 64        }
 65
 66        r
 67    }
 68    fn sin(x: f64) -> f64 {
 69        let mut r = x;
 70        let mut term = x;
 71        let mut n = 1;
 72
 73        for _ in 0..10 {
 74            term *= -(x * x) / ((2 * n) * (2 * n + 1)) as f64;
 75            r += term;
 76            n += 1;
 77        }
 78
 79        r
 80    }
 81
 82    // e^(self))
 83    pub fn exp(self) -> Self {
 84        Self {
 85            re: Self::cos(self.im), // TODO WIP review
 86            im: Self::sin(self.im),
 87        }
 88    }
 89    pub fn pow(self, k: u32) -> Self {
 90        let mut k = k;
 91        if k == 0 {
 92            return Self::one();
 93        }
 94        let mut base = self.clone();
 95        while k & 1 == 0 {
 96            base = base.clone() * base;
 97            k >>= 1;
 98        }
 99
100        if k == 1 {
101            return base;
102        }
103
104        let mut acc = base.clone();
105        while k > 1 {
106            k >>= 1;
107            base = base.clone() * base;
108            if k & 1 == 1 {
109                acc = acc * base.clone();
110            }
111        }
112        acc
113    }
114    pub fn modulus<const Q: u64>(self) -> Self {
115        let q: f64 = Q as f64;
116        let re = (self.re % q + q) % q;
117        let im = (self.im % q + q) % q;
118        Self { re, im }
119    }
120    pub fn modulus_centered<const Q: u64>(self) -> Self {
121        let re = modulus_centered_f64::<Q>(self.re);
122        let im = modulus_centered_f64::<Q>(self.im);
123        Self { re, im }
124    }
125}
126
127fn modulus_centered_f64<const Q: u64>(v: f64) -> f64 {
128    let q = Q as f64;
129    let mut res = v % q;
130    if res > q / 2.0 {
131        res = res - q;
132    }
133    res
134}
135
136impl<T: Default> Default for C<T> {
137    fn default() -> Self {
138        C {
139            re: T::default(),
140            im: T::default(),
141        }
142    }
143}
144
145impl<T> Add for C<T>
146where
147    T: Add<Output = T> + Copy,
148{
149    type Output = Self;
150    fn add(self, rhs: Self) -> Self::Output {
151        C {
152            re: self.re + rhs.re,
153            im: self.im + rhs.im,
154        }
155    }
156}
157
158impl<T> Sub for C<T>
159where
160    T: Sub<Output = T> + Copy,
161{
162    type Output = Self;
163    fn sub(self, rhs: Self) -> Self::Output {
164        C {
165            re: self.re - rhs.re,
166            im: self.im - rhs.im,
167        }
168    }
169}
170
171impl<T> Mul for C<T>
172where
173    T: Mul<Output = T> + Sub<Output = T> + Add<Output = T> + Copy,
174{
175    type Output = Self;
176    fn mul(self, rhs: Self) -> Self::Output {
177        C {
178            re: self.re * rhs.re - self.im * rhs.im,
179            im: self.re * rhs.im + self.im * rhs.re,
180        }
181    }
182}
183
184impl<T> Neg for C<T>
185where
186    T: Neg<Output = T> + Copy,
187{
188    type Output = Self;
189    fn neg(self) -> Self::Output {
190        C {
191            re: -self.re,
192            im: -self.im,
193        }
194    }
195}
196
197impl<T> Div for C<T>
198where
199    T: Copy
200        + Add<Output = T>
201        + Sub<Output = T>
202        + Mul<Output = T>
203        + Div<Output = T>
204        + Neg<Output = T>,
205{
206    type Output = Self;
207
208    fn div(self, rhs: Self) -> Self {
209        // (a+ib)/(c+id) = (ac + bd)/(c^2 + d^2) + i* (bc -ad)/(c^2 + d^2)
210        let den = rhs.re * rhs.re + rhs.im * rhs.im;
211        C {
212            re: (self.re * rhs.re + self.im * rhs.im) / den,
213            im: (self.im * rhs.re - self.re * rhs.im) / den,
214        }
215    }
216}
217impl<T> C<T>
218where
219    T: Neg<Output = T> + Copy,
220{
221    pub fn conj(&self) -> Self {
222        C {
223            re: self.re,
224            im: -self.im,
225        }
226    }
227}
228
229impl<T: Add<Output = T> + Default + Copy> std::iter::Sum for C<T> {
230    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
231        iter.fold(C::default(), |acc, x| acc + x)
232    }
233}
234
235impl C<f64> {
236    pub fn abs(&self) -> f64 {
237        (self.re * self.re + self.im * self.im).sqrt()
238    }
239}
240
241// poly mul with complex coefficients
242pub fn naive_poly_mul<const N: usize>(poly1: &Vec<C<f64>>, poly2: &Vec<C<f64>>) -> Vec<C<f64>> {
243    let mut result: Vec<C<f64>> = vec![C::<f64>::zero(); (N * 2) - 1];
244    for i in 0..N {
245        for j in 0..N {
246            result[i + j] = result[i + j] + poly1[i] * poly2[j];
247        }
248    }
249
250    // apply mod (X^N + 1))
251    // R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
252    // modulus_i128::<N>(&mut result);
253    // dbg!(&result);
254    // dbg!(R::<N>(array::from_fn(|i| result[i] as i64)).coeffs());
255
256    // R(array::from_fn(|i| result[i] as i64))
257    result
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263
264    #[test]
265    fn test_i() {
266        assert_eq!(C::i(), C::new(0.0, 1.0));
267    }
268
269    #[test]
270    fn test_add() {
271        let a = C::new(2.0, 3.0);
272        let b = C::new(1.0, 4.0);
273        assert_eq!(a + b, C::new(3.0, 7.0));
274    }
275
276    #[test]
277    fn test_sub() {
278        let a = C::new(5.0, 7.0);
279        let b = C::new(2.0, 3.0);
280        assert_eq!(a - b, C::new(3.0, 4.0));
281    }
282
283    #[test]
284    fn test_mult() {
285        let a = C::new(1.0, 2.0);
286        let b = C::new(3.0, 4.0);
287        assert_eq!(a * b, C::new(-5.0, 10.0));
288    }
289
290    #[test]
291    fn test_div() {
292        let a: C<f64> = C::new(1.0, 2.0);
293        let b: C<f64> = C::new(3.0, 4.0);
294        let r = a / b;
295        let expected = C::new(0.44, 0.08);
296        let epsilon = 1e-2;
297        assert!((r.re - expected.re).abs() < epsilon);
298        assert!((r.im - expected.im).abs() < epsilon);
299    }
300
301    #[test]
302    fn test_conj() {
303        let a = C::new(3.0, -4.0);
304        assert_eq!(a.conj(), C::new(3.0, 4.0));
305        assert_eq!(a.conj().conj(), a);
306    }
307
308    #[test]
309    fn test_neg() {
310        let a = C::new(1.0, -2.0);
311        assert_eq!(-a, C::new(-1.0, 2.0));
312    }
313
314    #[test]
315    fn test_abs() {
316        let a = C::new(3.0, 4.0);
317        assert!((a.abs() - 5.0).abs() < 1e-10);
318    }
319
320    #[test]
321    fn test_exp() {
322        // let a = C::new(3.0, 4.0);
323        let pi = C::<f64>::from(std::f64::consts::PI);
324        let n = 4;
325        let a = ((C::<f64>::from(2f64) * pi * C::<f64>::i()) / C::<f64>::new(n as f64, 0f64)).exp();
326        dbg!(&a);
327        assert_eq!(a.exp(), a.exp());
328    }
329}