scirs2_fft/
fftw_backend.rs1use 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
14pub fn rfft_fftw(input: &ArrayView1<f64>) -> FFTResult<Array1<Complex<f64>>> {
27 let n = input.len();
28
29 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 fftw_plan_cache::execute_r2c(&mut input_vec, &mut output)?;
36
37 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
43pub fn fft_fftw(input: &ArrayView1<Complex<f64>>) -> FFTResult<Array1<Complex<f64>>> {
47 let n = input.len();
48
49 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 fftw_plan_cache::execute_c2c_forward(&mut input_fftw, &mut output)?;
55
56 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
62pub fn ifft_fftw(input: &ArrayView1<Complex<f64>>) -> FFTResult<Array1<Complex<f64>>> {
64 let n = input.len();
65
66 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 fftw_plan_cache::execute_c2c_backward(&mut input_fftw, &mut output)?;
72
73 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
83pub fn irfft_fftw(input: &ArrayView1<Complex<f64>>, n: usize) -> FFTResult<Array1<f64>> {
87 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 fftw_plan_cache::execute_c2r(&mut input_fftw, &mut output, n)?;
93
94 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
101pub fn fft2_fftw(input: &ArrayView2<Complex<f64>>) -> FFTResult<Array2<Complex<f64>>> {
107 let (rows, cols) = input.dim();
108 let n = rows * cols;
109
110 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 fftw_plan_cache::execute_c2c_2d_forward(&mut input_fftw, &mut output, rows, cols)?;
116
117 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
124pub fn ifft2_fftw(input: &ArrayView2<Complex<f64>>) -> FFTResult<Array2<Complex<f64>>> {
126 let (rows, cols) = input.dim();
127 let n = rows * cols;
128
129 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 fftw_plan_cache::execute_c2c_2d_backward(&mut input_fftw, &mut output, rows, cols)?;
135
136 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
147pub fn rfft2_fftw(input: &ArrayView2<f64>) -> FFTResult<Array2<Complex<f64>>> {
149 let (rows, cols) = input.dim();
150
151 let mut input_vec: Vec<f64> = input.iter().cloned().collect();
153
154 let out_cols = cols / 2 + 1;
156 let mut output = vec![c64::default(); rows * out_cols];
157
158 fftw_plan_cache::execute_r2c_2d(&mut input_vec, &mut output, rows, cols)?;
160
161 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
168pub 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 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 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 fftw_plan_cache::execute_c2r_2d(&mut input_fftw, &mut output, rows, cols)?;
195
196 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
204pub 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 fftw_plan_cache::execute_dct2(&mut input_vec, &mut output)?;
217
218 Ok(Array1::from_vec(output))
219}
220
221pub 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 fftw_plan_cache::execute_idct2(&mut input_vec, &mut output)?;
230
231 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
238pub 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 fftw_plan_cache::execute_dst2(&mut input_vec, &mut output)?;
246
247 Ok(Array1::from_vec(output))
248}
249
250pub 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 fftw_plan_cache::execute_idst2(&mut input_vec, &mut output)?;
258
259 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}