Skip to main content

embedded_dsp/
transform.rs

1//! Fast Fourier Transform (FFT), Real FFT (RFFT), Discrete Cosine Transform (DCT-IV), and Bit Reversal functions.
2
3#[allow(unused_imports)]
4use crate::math::FloatMath;
5use crate::types::*;
6
7/// Bit reversal function for interleaved complex array of size `2 * n`.
8pub fn bit_reversal(data: &mut [f32], n: usize) {
9    let mut j = 0;
10    for i in 0..n {
11        if i < j {
12            data.swap(2 * i, 2 * j);
13            data.swap(2 * i + 1, 2 * j + 1);
14        }
15        let mut m = n >> 1;
16        while m >= 1 && j >= m {
17            j -= m;
18            m >>= 1;
19        }
20        j += m;
21    }
22}
23
24/// In-place Complex FFT for floating point 32-bit (`f32`).
25/// `data` is interleaved complex array of size `2 * n` (`[re0, im0, re1, im1, ...]`).
26/// `ifft_flag`: 0 for forward FFT, 1 for inverse FFT (IFFT).
27/// `bit_reverse_flag`: 1 to enable bit reversal, 0 to disable.
28pub fn cfft_f32(data: &mut [f32], n: usize, ifft_flag: u8, bit_reverse_flag: u8) {
29    if n < 2 || (n & (n - 1)) != 0 {
30        return;
31    }
32
33    if bit_reverse_flag != 0 {
34        bit_reversal(data, n);
35    }
36
37    let mut len = 2;
38    while len <= n {
39        let half_len = len / 2;
40        let angle =
41            (if ifft_flag != 0 { 2.0 } else { -2.0 }) * core::f32::consts::PI / (len as f32);
42        let w_step_re = angle.cos();
43        let w_step_im = angle.sin();
44
45        let mut i = 0;
46        while i < n {
47            let mut w_re = 1.0f32;
48            let mut w_im = 0.0f32;
49
50            for j in 0..half_len {
51                let u_idx = 2 * (i + j);
52                let v_idx = 2 * (i + j + half_len);
53
54                let u_re = data[u_idx];
55                let u_im = data[u_idx + 1];
56
57                let v_re = data[v_idx];
58                let v_im = data[v_idx + 1];
59
60                let t_re = v_re * w_re - v_im * w_im;
61                let t_im = v_re * w_im + v_im * w_re;
62
63                data[u_idx] = u_re + t_re;
64                data[u_idx + 1] = u_im + t_im;
65
66                data[v_idx] = u_re - t_re;
67                data[v_idx + 1] = u_im - t_im;
68
69                let next_w_re = w_re * w_step_re - w_im * w_step_im;
70                let next_w_im = w_re * w_step_im + w_im * w_step_re;
71                w_re = next_w_re;
72                w_im = next_w_im;
73            }
74            i += len;
75        }
76        len <<= 1;
77    }
78
79    if ifft_flag != 0 {
80        let norm = 1.0 / (n as f32);
81        for i in 0..(2 * n) {
82            data[i] *= norm;
83        }
84    }
85}
86
87/// In-place Complex FFT for Q31 fixed-point.
88pub fn cfft_q31(data: &mut [q31], n: usize, ifft_flag: u8, _bit_reverse_flag: u8) {
89    if n < 2 {
90        return;
91    }
92    // Convert to f32 scratch, run cfft_f32, convert back
93    let mut scratch = [0.0f32; 1024];
94    let total = 2 * n;
95    if total > scratch.len() {
96        return;
97    }
98
99    for i in 0..total {
100        scratch[i] = data[i] as f32 / 2147483648.0;
101    }
102    cfft_f32(&mut scratch[..total], n, ifft_flag, 1);
103    for i in 0..total {
104        data[i] = (scratch[i] * 2147483648.0).clamp(-2147483648.0, 2147483647.0) as q31;
105    }
106}
107
108/// In-place Complex FFT for Q15 fixed-point.
109pub fn cfft_q15(data: &mut [q15], n: usize, ifft_flag: u8, _bit_reverse_flag: u8) {
110    if n < 2 {
111        return;
112    }
113    let mut scratch = [0.0f32; 1024];
114    let total = 2 * n;
115    if total > scratch.len() {
116        return;
117    }
118
119    for i in 0..total {
120        scratch[i] = data[i] as f32 / 32768.0;
121    }
122    cfft_f32(&mut scratch[..total], n, ifft_flag, 1);
123    for i in 0..total {
124        data[i] = (scratch[i] * 32768.0).clamp(-32768.0, 32767.0) as q15;
125    }
126}
127
128/// Real FFT for floating point 32-bit (`f32`).
129/// `src` has `n` real samples. `dst` receives `2 * n` complex outputs.
130pub fn rfft_f32(src: &[f32], dst: &mut [f32], n: usize, ifft_flag: u8) {
131    let len = src.len().min(n);
132    let mut c_data = [0.0f32; 1024];
133    if 2 * len > c_data.len() || dst.len() < 2 * len {
134        return;
135    }
136
137    for i in 0..len {
138        c_data[2 * i] = src[i];
139        c_data[2 * i + 1] = 0.0;
140    }
141
142    cfft_f32(&mut c_data[..2 * len], len, ifft_flag, 1);
143    dst[..2 * len].copy_from_slice(&c_data[..2 * len]);
144}
145
146/// Real FFT for Q31 fixed-point.
147pub fn rfft_q31(src: &[q31], dst: &mut [q31], n: usize, ifft_flag: u8) {
148    let len = src.len().min(n);
149    let mut c_data = [0; 1024];
150    if 2 * len > c_data.len() || dst.len() < 2 * len {
151        return;
152    }
153
154    for i in 0..len {
155        c_data[2 * i] = src[i];
156        c_data[2 * i + 1] = 0;
157    }
158    cfft_q31(&mut c_data[..2 * len], len, ifft_flag, 1);
159    dst[..2 * len].copy_from_slice(&c_data[..2 * len]);
160}
161
162/// Real FFT for Q15 fixed-point.
163pub fn rfft_q15(src: &[q15], dst: &mut [q15], n: usize, ifft_flag: u8) {
164    let len = src.len().min(n);
165    let mut c_data = [0; 1024];
166    if 2 * len > c_data.len() || dst.len() < 2 * len {
167        return;
168    }
169
170    for i in 0..len {
171        c_data[2 * i] = src[i];
172        c_data[2 * i + 1] = 0;
173    }
174    cfft_q15(&mut c_data[..2 * len], len, ifft_flag, 1);
175    dst[..2 * len].copy_from_slice(&c_data[..2 * len]);
176}
177
178/// Discrete Cosine Transform Type IV (DCT-IV) for f32.
179pub fn dct4_f32(src: &[f32], dst: &mut [f32], n: usize) {
180    let len = src.len().min(dst.len()).min(n);
181    let pi_over_n = core::f32::consts::PI / (len as f32);
182
183    for k in 0..len {
184        let mut sum = 0.0f32;
185        let k_factor = (k as f32 + 0.5) * pi_over_n;
186        for n_idx in 0..len {
187            let angle = (n_idx as f32 + 0.5) * k_factor;
188            sum += src[n_idx] * angle.cos();
189        }
190        let norm = (2.0 / len as f32).sqrt();
191        dst[k] = sum * norm;
192    }
193}