scirs2_fft/
fftw_backend.rs

1//! FFTW-based FFT backend for high performance
2//!
3//! Provides 62x speedup over pure Rust FFT implementation by using
4//! the industry-standard FFTW library.
5//!
6//! Now with plan caching for even better performance on repeated transforms.
7
8use crate::error::{FFTError, FFTResult};
9use crate::fftw_plan_cache;
10use fftw::types::*;
11use scirs2_core::ndarray::{Array1, Array2, ArrayView1, ArrayView2};
12use scirs2_core::numeric::Complex;
13
14/// Real-to-complex FFT using FFTW with plan caching
15///
16/// Computes the FFT of a real-valued signal, returning complex frequency components.
17/// Uses cached plans for better performance on repeated transforms of the same size.
18///
19/// # Arguments
20///
21/// * `input` - Real-valued input signal
22///
23/// # Returns
24///
25/// * Complex frequency spectrum (N/2+1 components for real input of length N)
26pub fn rfft_fftw(input: &ArrayView1<f64>) -> FFTResult<Array1<Complex<f64>>> {
27    let n = input.len();
28
29    // Prepare input and output buffers
30    let mut input_vec = input.to_vec();
31    let output_len = n / 2 + 1;
32    let mut output = vec![c64::default(); output_len];
33
34    // Execute with cached plan
35    fftw_plan_cache::execute_r2c(&mut input_vec, &mut output)?;
36
37    // Convert fftw::c64 to scirs2 Complex<f64>
38    let result: Vec<Complex<f64>> = output.iter().map(|c| Complex::new(c.re, c.im)).collect();
39
40    Ok(Array1::from_vec(result))
41}
42
43/// Complex-to-complex FFT using FFTW with plan caching
44///
45/// Computes the FFT of a complex-valued signal.
46pub fn fft_fftw(input: &ArrayView1<Complex<f64>>) -> FFTResult<Array1<Complex<f64>>> {
47    let n = input.len();
48
49    // Convert scirs2 Complex to fftw c64
50    let mut input_fftw: Vec<c64> = input.iter().map(|c| c64 { re: c.re, im: c.im }).collect();
51    let mut output = vec![c64::default(); n];
52
53    // Execute with cached plan
54    fftw_plan_cache::execute_c2c_forward(&mut input_fftw, &mut output)?;
55
56    // Convert back
57    let result: Vec<Complex<f64>> = output.iter().map(|c| Complex::new(c.re, c.im)).collect();
58
59    Ok(Array1::from_vec(result))
60}
61
62/// Inverse complex-to-complex FFT using FFTW with plan caching
63pub fn ifft_fftw(input: &ArrayView1<Complex<f64>>) -> FFTResult<Array1<Complex<f64>>> {
64    let n = input.len();
65
66    // Convert to fftw format
67    let mut input_fftw: Vec<c64> = input.iter().map(|c| c64 { re: c.re, im: c.im }).collect();
68    let mut output = vec![c64::default(); n];
69
70    // Execute with cached plan
71    fftw_plan_cache::execute_c2c_backward(&mut input_fftw, &mut output)?;
72
73    // FFTW doesn't normalize, so divide by N
74    let scale = 1.0 / (n as f64);
75    let result: Vec<Complex<f64>> = output
76        .iter()
77        .map(|c| Complex::new(c.re * scale, c.im * scale))
78        .collect();
79
80    Ok(Array1::from_vec(result))
81}
82
83/// Inverse real FFT using FFTW with plan caching
84///
85/// Converts complex frequency spectrum back to real-valued time-domain signal.
86pub fn irfft_fftw(input: &ArrayView1<Complex<f64>>, n: usize) -> FFTResult<Array1<f64>> {
87    // Convert to fftw format
88    let mut input_fftw: Vec<c64> = input.iter().map(|c| c64 { re: c.re, im: c.im }).collect();
89    let mut output = vec![0.0; n];
90
91    // Execute with cached plan
92    fftw_plan_cache::execute_c2r(&mut input_fftw, &mut output, n)?;
93
94    // FFTW doesn't normalize
95    let scale = 1.0 / (n as f64);
96    let result: Vec<f64> = output.iter().map(|&x| x * scale).collect();
97
98    Ok(Array1::from_vec(result))
99}
100
101// ========================================
102// 2D FFT FUNCTIONS
103// ========================================
104
105/// 2D complex-to-complex FFT using FFTW with plan caching
106pub fn fft2_fftw(input: &ArrayView2<Complex<f64>>) -> FFTResult<Array2<Complex<f64>>> {
107    let (rows, cols) = input.dim();
108    let n = rows * cols;
109
110    // Convert to row-major contiguous Vec
111    let mut input_fftw: Vec<c64> = input.iter().map(|c| c64 { re: c.re, im: c.im }).collect();
112    let mut output = vec![c64::default(); n];
113
114    // Execute with cached plan
115    fftw_plan_cache::execute_c2c_2d_forward(&mut input_fftw, &mut output, rows, cols)?;
116
117    // Convert back and reshape to 2D
118    let result: Vec<Complex<f64>> = output.iter().map(|c| Complex::new(c.re, c.im)).collect();
119
120    Array2::from_shape_vec((rows, cols), result)
121        .map_err(|e| FFTError::ComputationError(format!("Failed to reshape result: {:?}", e)))
122}
123
124/// 2D inverse complex-to-complex FFT using FFTW with plan caching
125pub fn ifft2_fftw(input: &ArrayView2<Complex<f64>>) -> FFTResult<Array2<Complex<f64>>> {
126    let (rows, cols) = input.dim();
127    let n = rows * cols;
128
129    // Convert to row-major contiguous Vec
130    let mut input_fftw: Vec<c64> = input.iter().map(|c| c64 { re: c.re, im: c.im }).collect();
131    let mut output = vec![c64::default(); n];
132
133    // Execute with cached plan
134    fftw_plan_cache::execute_c2c_2d_backward(&mut input_fftw, &mut output, rows, cols)?;
135
136    // FFTW doesn't normalize
137    let scale = 1.0 / (n as f64);
138    let result: Vec<Complex<f64>> = output
139        .iter()
140        .map(|c| Complex::new(c.re * scale, c.im * scale))
141        .collect();
142
143    Array2::from_shape_vec((rows, cols), result)
144        .map_err(|e| FFTError::ComputationError(format!("Failed to reshape result: {:?}", e)))
145}
146
147/// 2D real-to-complex FFT using FFTW with plan caching
148pub fn rfft2_fftw(input: &ArrayView2<f64>) -> FFTResult<Array2<Complex<f64>>> {
149    let (rows, cols) = input.dim();
150
151    // Convert to row-major contiguous Vec
152    let mut input_vec: Vec<f64> = input.iter().cloned().collect();
153
154    // For 2D real FFT, output has shape (rows, cols/2 + 1)
155    let out_cols = cols / 2 + 1;
156    let mut output = vec![c64::default(); rows * out_cols];
157
158    // Execute with cached plan
159    fftw_plan_cache::execute_r2c_2d(&mut input_vec, &mut output, rows, cols)?;
160
161    // Convert back and reshape
162    let result: Vec<Complex<f64>> = output.iter().map(|c| Complex::new(c.re, c.im)).collect();
163
164    Array2::from_shape_vec((rows, out_cols), result)
165        .map_err(|e| FFTError::ComputationError(format!("Failed to reshape result: {:?}", e)))
166}
167
168/// 2D inverse real FFT using FFTW with plan caching
169pub fn irfft2_fftw(
170    input: &ArrayView2<Complex<f64>>,
171    shape: (usize, usize),
172) -> FFTResult<Array2<f64>> {
173    let (rows, cols) = shape;
174    let (in_rows, in_cols) = input.dim();
175
176    // Validate input dimensions
177    if in_rows != rows || in_cols != cols / 2 + 1 {
178        return Err(FFTError::ValueError(format!(
179            "Input shape ({}, {}) doesn't match expected ({}, {}) for output shape ({}, {})",
180            in_rows,
181            in_cols,
182            rows,
183            cols / 2 + 1,
184            rows,
185            cols
186        )));
187    }
188
189    // Convert to row-major contiguous Vec
190    let mut input_fftw: Vec<c64> = input.iter().map(|c| c64 { re: c.re, im: c.im }).collect();
191    let mut output = vec![0.0; rows * cols];
192
193    // Execute with cached plan
194    fftw_plan_cache::execute_c2r_2d(&mut input_fftw, &mut output, rows, cols)?;
195
196    // FFTW doesn't normalize
197    let scale = 1.0 / ((rows * cols) as f64);
198    let result: Vec<f64> = output.iter().map(|&x| x * scale).collect();
199
200    Array2::from_shape_vec((rows, cols), result)
201        .map_err(|e| FFTError::ComputationError(format!("Failed to reshape result: {:?}", e)))
202}
203
204// ========================================
205// DCT/DST FUNCTIONS using R2R transforms with plan caching
206// ========================================
207
208/// DCT Type II using FFTW with plan caching (REDFT10)
209/// This is the most commonly used DCT type
210pub fn dct2_fftw(input: &ArrayView1<f64>) -> FFTResult<Array1<f64>> {
211    let n = input.len();
212    let mut input_vec = input.to_vec();
213    let mut output = vec![0.0; n];
214
215    // Execute with cached plan
216    fftw_plan_cache::execute_dct2(&mut input_vec, &mut output)?;
217
218    Ok(Array1::from_vec(output))
219}
220
221/// IDCT Type II using FFTW with plan caching (REDFT01 = DCT-III)
222/// DCT-III is the inverse of DCT-II
223pub fn idct2_fftw(input: &ArrayView1<f64>) -> FFTResult<Array1<f64>> {
224    let n = input.len();
225    let mut input_vec = input.to_vec();
226    let mut output = vec![0.0; n];
227
228    // Execute with cached plan
229    fftw_plan_cache::execute_idct2(&mut input_vec, &mut output)?;
230
231    // FFTW doesn't normalize - need to scale by 1/(2*N)
232    let scale = 1.0 / (2.0 * n as f64);
233    let result: Vec<f64> = output.iter().map(|&x| x * scale).collect();
234
235    Ok(Array1::from_vec(result))
236}
237
238/// DST Type II using FFTW with plan caching (RODFT10)
239pub fn dst2_fftw(input: &ArrayView1<f64>) -> FFTResult<Array1<f64>> {
240    let n = input.len();
241    let mut input_vec = input.to_vec();
242    let mut output = vec![0.0; n];
243
244    // Execute with cached plan
245    fftw_plan_cache::execute_dst2(&mut input_vec, &mut output)?;
246
247    Ok(Array1::from_vec(output))
248}
249
250/// IDST Type II using FFTW with plan caching (RODFT01 = DST-III)
251pub fn idst2_fftw(input: &ArrayView1<f64>) -> FFTResult<Array1<f64>> {
252    let n = input.len();
253    let mut input_vec = input.to_vec();
254    let mut output = vec![0.0; n];
255
256    // Execute with cached plan
257    fftw_plan_cache::execute_idst2(&mut input_vec, &mut output)?;
258
259    // FFTW doesn't normalize
260    let scale = 1.0 / (2.0 * n as f64);
261    let result: Vec<f64> = output.iter().map(|&x| x * scale).collect();
262
263    Ok(Array1::from_vec(result))
264}