makepad_math/
complex.rs

1
2// very basic complex types just enough to write an fft below
3
4#[derive(Clone, Copy)]
5pub struct ComplexF32 {pub re: f32, pub im: f32}
6
7#[derive(Clone, Copy)]
8pub struct ComplexF64 {pub re: f64, pub im: f64}
9
10impl ComplexF32{
11    pub fn magnitude(self)->f32{(self.re*self.re + self.im *self.im).sqrt()}
12}
13
14pub fn cf64(re: f64, im: f64) -> ComplexF64 {ComplexF64 {re, im}}
15pub fn cf32(re: f32, im: f32) -> ComplexF32 {ComplexF32 {re, im}}
16
17impl From<ComplexF32> for ComplexF64 {
18    fn from(v: ComplexF32) -> Self {
19        cf64(v.re as f64, v.im as f64)
20    }
21}
22
23impl From<ComplexF64> for ComplexF32 {
24    fn from(v: ComplexF64) -> Self {
25        cf32(v.re as f32, v.im as f32)
26    }
27}
28
29impl std::ops::Mul<ComplexF64> for ComplexF64 {
30    type Output = ComplexF64;
31    fn mul(self, rhs: ComplexF64) -> ComplexF64 {
32        cf64(
33            self.re * rhs.re - self.im * rhs.im,
34            self.re * rhs.im + self.im * rhs.re
35        )
36    }
37}
38
39impl std::ops::Add<ComplexF64> for ComplexF64 {
40    type Output = ComplexF64;
41    fn add(self, rhs: ComplexF64) -> ComplexF64 {
42        cf64(self.re + rhs.re, self.im + rhs.im)
43    }
44}
45
46impl std::ops::Sub<ComplexF64> for ComplexF64 {
47    type Output = ComplexF64;
48    fn sub(self, rhs: ComplexF64) -> ComplexF64 {
49        cf64(self.re - rhs.re, self.im - rhs.im)
50    }
51}
52
53// FFT algo, ported from https://github.com/rshuston/FFT-C/ rewritten with a few Rust types.
54
55fn fft_f32_recursive_pow2_inner(data: &mut [ComplexF32], scratch: &mut [ComplexF32], n: usize, theta_pi: f64, stride: usize) {
56    if stride < n {
57        let stride2 = stride * 2;
58        fft_f32_recursive_pow2_inner(scratch, data, n, theta_pi, stride2);
59        fft_f32_recursive_pow2_inner(&mut scratch[stride..], &mut data[stride..], n, theta_pi, stride2);
60        
61        let theta = (stride2 as f64 * theta_pi) / n as f64;
62        let wn = cf64(theta.cos(), theta.sin());
63        let mut wnk = cf64(1.0, 0.0);
64
65        for k in (0..n).step_by(stride2) {
66            let kd2 = k >> 1;
67            let kpnd2 = (k + n) >> 1;
68            
69            let u: ComplexF64 = scratch[k].into();
70            let t: ComplexF64 = wnk * scratch[k + stride].into();
71            
72            data[kd2] = (u + t).into();
73            data[kpnd2] = (u - t).into();
74            
75            wnk = wnk * wn;
76        }
77    }
78}
79
80use std::f64::consts::PI;
81
82pub fn fft_f32_recursive_pow2_forward(data: &mut [ComplexF32], scratch: &mut [ComplexF32]) {
83    fft_f32_recursive_pow2(data, scratch, -PI)
84}
85
86pub fn fft_f32_recursive_pow2_inverse(data: &mut [ComplexF32], scratch: &mut [ComplexF32]) {
87    fft_f32_recursive_pow2(data, scratch, PI)
88}
89
90fn fft_f32_recursive_pow2(data: &mut [ComplexF32], scratch: &mut [ComplexF32], theta_pi: f64) {
91    let n = data.len();
92    if data.len() != scratch.len() {panic!()}
93    fn is_power_of_2(n: usize)->bool{n != 0 && (!(n & (n - 1))) != 0 }
94    if !is_power_of_2(n){ // check power of two
95        panic!("fft data length must be power of 2");
96    };
97    scratch.copy_from_slice(data);
98    fft_f32_recursive_pow2_inner(data, scratch, n, theta_pi, 1);
99}
makepad_math/complex.rs

makepad_math/
complex.rs
