Complex numbers and FFTs

Cargo.toml

@@ -1,7 +1,9 @@
 [workspace]
 members = [
   "hcpl_algebra",
+  "hcpl_complex",
   "hcpl_divide_and_conquer_dp",
+  "hcpl_fft",
   "hcpl_fwht",
   "hcpl_integer",
   "hcpl_io",

hcpl_complex/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "hcpl_complex"
+version = "0.1.0"
+edition = "2021"
+repository = "https://github.com/THE-nio/hcpl"
+license = "MIT"
+
+[dependencies]
+hcpl_algebra = { version = "0.1.0", path = "../hcpl_algebra" }
+hcpl_number_theory = { path = "../hcpl_number_theory" }

hcpl_complex/src/lib.rs

@@ -0,0 +1,256 @@
+#[derive(Clone, Copy, Debug, Default)]
+pub struct Complex {
+    pub re: f64,
+    pub im: f64,
+}
+
+impl Complex {
+    pub const ZERO: Complex = Complex { re: 0., im: 0. };
+    pub const ONE: Complex = Complex { re: 1., im: 0. };
+    pub const I: Complex = Complex { re: 0., im: 1. };
+
+    pub fn powi(mut self, mut rhs: u32) -> Self {
+        let mut result = Self::ONE;
+
+        while rhs > 0 {
+            if (rhs & 1) != 0 {
+                result *= self;
+            }
+
+            self *= self;
+            rhs >>= 1;
+        }
+
+        result
+    }
+
+    pub fn cis(angle: f64) -> Self {
+        Self {
+            re: angle.cos(),
+            im: angle.sin(),
+        }
+    }
+
+    pub fn conj(self) -> Self {
+        Self {
+            re: self.re,
+            im: -self.im,
+        }
+    }
+
+    pub fn squared_norm(self) -> f64 {
+        self.re * self.re + self.im * self.im
+    }
+
+    pub fn inv(self) -> Self {
+        Self {
+            re: self.re / self.squared_norm(),
+            im: -self.im / self.squared_norm(),
+        }
+    }
+}
+
+impl std::fmt::Display for Complex {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_fmt(format_args!("{} + i {}", self.re, self.im))
+    }
+}
+
+impl From<f64> for Complex {
+    fn from(value: f64) -> Self {
+        Self { re: value, im: 0. }
+    }
+}
+
+impl From<usize> for Complex {
+    fn from(value: usize) -> Self {
+        Self {
+            re: value as f64,
+            im: 0.,
+        }
+    }
+}
+
+impl std::ops::Add<Complex> for Complex {
+    type Output = Complex;
+
+    fn add(self, rhs: Complex) -> Self::Output {
+        Self {
+            re: self.re + rhs.re,
+            im: self.im + rhs.im,
+        }
+    }
+}
+
+impl std::ops::Add<f64> for Complex {
+    type Output = Complex;
+
+    fn add(self, rhs: f64) -> Self::Output {
+        self + Complex::from(rhs)
+    }
+}
+
+impl std::ops::Neg for Complex {
+    type Output = Complex;
+
+    fn neg(self) -> Self::Output {
+        Self {
+            re: -self.re,
+            im: -self.im,
+        }
+    }
+}
+
+impl std::ops::Sub<Complex> for Complex {
+    type Output = Complex;
+
+    fn sub(self, rhs: Complex) -> Self::Output {
+        self + -rhs
+    }
+}
+
+impl std::ops::Sub<f64> for Complex {
+    type Output = Complex;
+
+    fn sub(self, rhs: f64) -> Self::Output {
+        self + -rhs
+    }
+}
+
+impl std::ops::Mul<Complex> for Complex {
+    type Output = Complex;
+
+    fn mul(self, rhs: Complex) -> Self::Output {
+        Self {
+            re: self.re * rhs.re - self.im * rhs.im,
+            im: self.re * rhs.im + self.im * rhs.re,
+        }
+    }
+}
+
+impl std::ops::Mul<f64> for Complex {
+    type Output = Complex;
+
+    fn mul(self, rhs: f64) -> Self::Output {
+        Self {
+            re: self.re * rhs,
+            im: self.im * rhs,
+        }
+    }
+}
+
+impl std::ops::Div<Complex> for Complex {
+    type Output = Complex;
+
+    fn div(self, rhs: Complex) -> Self::Output {
+        self * rhs.inv()
+    }
+}
+
+impl std::ops::Div<f64> for Complex {
+    type Output = Complex;
+
+    fn div(self, rhs: f64) -> Self::Output {
+        Self {
+            re: self.re / rhs,
+            im: self.im / rhs,
+        }
+    }
+}
+
+impl std::ops::AddAssign<Complex> for Complex {
+    fn add_assign(&mut self, rhs: Complex) {
+        *self = *self + rhs;
+    }
+}
+
+impl std::ops::AddAssign<f64> for Complex {
+    fn add_assign(&mut self, rhs: f64) {
+        *self = *self + rhs;
+    }
+}
+
+impl std::ops::SubAssign<Complex> for Complex {
+    fn sub_assign(&mut self, rhs: Complex) {
+        *self = *self - rhs;
+    }
+}
+
+impl std::ops::SubAssign<f64> for Complex {
+    fn sub_assign(&mut self, rhs: f64) {
+        *self = *self - rhs;
+    }
+}
+
+impl std::ops::MulAssign<Complex> for Complex {
+    fn mul_assign(&mut self, rhs: Complex) {
+        *self = *self * rhs;
+    }
+}
+
+impl std::ops::MulAssign<f64> for Complex {
+    fn mul_assign(&mut self, rhs: f64) {
+        *self = *self * rhs;
+    }
+}
+
+impl std::ops::DivAssign<Complex> for Complex {
+    fn div_assign(&mut self, rhs: Complex) {
+        *self = *self / rhs;
+    }
+}
+
+impl std::ops::DivAssign<f64> for Complex {
+    fn div_assign(&mut self, rhs: f64) {
+        *self = *self / rhs;
+    }
+}
+
+impl std::iter::Sum<Complex> for Complex {
+    fn sum<I: Iterator<Item = Complex>>(iter: I) -> Self {
+        let mut result = Complex::ZERO;
+
+        for item in iter {
+            result += item;
+        }
+
+        result
+    }
+}
+
+impl std::iter::Product<Complex> for Complex {
+    fn product<I: Iterator<Item = Complex>>(iter: I) -> Self {
+        let mut result = Complex::ONE;
+
+        for item in iter {
+            result *= item;
+        }
+
+        result
+    }
+}
+
+impl hcpl_algebra::monoid::AdditiveIdentity for Complex {
+    const VALUE: Self = Complex::ZERO;
+}
+impl hcpl_algebra::monoid::MultiplicativeIdentity for Complex {
+    const VALUE: Self = Complex::ONE;
+}
+
+impl hcpl_number_theory::roots::TryNthRootOfUnity for Complex {
+    type Error = std::convert::Infallible;
+
+    fn try_nth_root_of_unity(n: usize) -> Result<Self, Self::Error>
+    where
+        Self: Sized,
+    {
+        Ok(Self::cis(2. * std::f64::consts::PI / n as f64))
+    }
+
+    fn try_nth_root_of_unity_inv(n: usize) -> Result<Self, Self::Error>
+    where
+        Self: Sized,
+    {
+        Ok(Self::cis(-2. * std::f64::consts::PI / n as f64))
+    }
+}

hcpl_fft/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "hcpl_fft"
+version = "0.1.0"
+edition = "2021"
+repository = "https://github.com/THE-nio/hcpl"
+license = "MIT"
+
+[dependencies]
+hcpl_algebra = { path = "../hcpl_algebra" }
+hcpl_number_theory = { path = "../hcpl_number_theory" }
+hcpl_complex = { path = "../hcpl_complex", optional = true }
+
+[features]
+real = [ "dep:hcpl_complex" ]

hcpl_fft/src/lib.rs

@@ -0,0 +1,98 @@
+#[cfg(feature = "real")]
+pub mod real;
+
+mod stack_stack;
+
+use hcpl_algebra::monoid::MultiplicativeIdentity;
+use hcpl_number_theory::roots::TryNthRootOfUnity;
+use std::fmt::Debug;
+
+/// Performs the bit-reversal permutation on `vec`
+pub fn bit_reversal<T>(vec: &mut [T]) {
+    let n = vec.len();
+
+    let mut j = 0;
+    for i in 1..n {
+        let mut bit = n >> 1;
+        while (j & bit) != 0 {
+            j ^= bit;
+            bit >>= 1;
+        }
+        j ^= bit;
+
+        if i < j {
+            vec.swap(i, j);
+        }
+    }
+}
+
+fn get_roots<T>(mut n: usize, inv: bool) -> stack_stack::StackStack<T>
+where
+    T: TryNthRootOfUnity,
+    <T as TryNthRootOfUnity>::Error: Debug,
+{
+    let mut roots = stack_stack::StackStack::new();
+
+    let mut last: T = if inv {
+        TryNthRootOfUnity::try_nth_root_of_unity_inv(n)
+    } else {
+        TryNthRootOfUnity::try_nth_root_of_unity(n)
+    }
+    .unwrap();
+    while n >= 2 {
+        roots.push(last);
+        n /= 2;
+        last = last * last;
+    }
+
+    roots
+}
+
+/// Performs the in-place Fast Fourier Transform on the slice `vec`, whose lenght must be a power of two. DOES NOT
+/// PERFORM NORMALISATION
+pub fn fft<T>(vec: &mut [T], inv: bool)
+where
+    T: TryNthRootOfUnity,
+    <T as TryNthRootOfUnity>::Error: Debug,
+{
+    bit_reversal(vec);
+
+    let n = vec.len();
+
+    let mut roots = get_roots(n, inv);
+    let mut width = 2;
+    while width <= n {
+        let w_d = roots.pop().unwrap();
+
+        for i in (0..n).step_by(width) {
+            let mut w = <T as MultiplicativeIdentity>::VALUE;
+            for j in 0..width / 2 {
+                let l = i + j;
+                let r = i + j + width / 2;
+
+                (vec[l], vec[r]) = (vec[l] + w * vec[r], vec[l] - w * vec[r]);
+                w = w * w_d;
+            }
+        }
+
+        width *= 2;
+    }
+}
+
+/// Performs the in-place discrete convolution of the two vectors `a` and `b`. DOES NOT PERFORM
+/// NORMALISATION
+///
+/// `a` will contain the convolution of `a` and `b`, and `b` will contain the discrete
+/// Fourier transform of `b`.
+pub fn convolution<T>(a: &mut [T], b: &mut [T])
+where
+    T: TryNthRootOfUnity,
+    <T as TryNthRootOfUnity>::Error: Debug,
+{
+    fft(a, false);
+    fft(b, false);
+    for (x, y) in a.iter_mut().zip(b.iter().copied()) {
+        *x = *x * y;
+    }
+    fft(a, true);
+}