rusty_psf/fft/
transform.rs

1use ndarray::{Array2, ArrayView2, s};
2use num_complex::Complex64;
3use rustfft::FftPlanner;
4
5/// Perform 2D FFT on a complex-valued array.
6/// 
7/// This function performs a 2D Fast Fourier Transform using row-wise and column-wise
8/// 1D FFTs. The input array is processed in-place to minimize memory usage.
9/// 
10/// # Arguments
11/// * `input` - Complex-valued 2D array to transform
12/// 
13/// # Returns
14/// A new 2D array containing the Fourier transform
15pub fn fft2d(input: &ArrayView2<Complex64>) -> Array2<Complex64> {
16    let (rows, cols) = input.dim();
17    let mut output = input.to_owned();
18    let mut planner = FftPlanner::new();
19    
20    // Row-wise FFT
21    let row_fft = planner.plan_fft_forward(cols);
22    for mut row in output.rows_mut() {
23        // Convert row to Vec to ensure contiguous memory
24        let mut row_vec = row.to_vec();
25        row_fft.process(&mut row_vec);
26        // Copy back to row
27        row.iter_mut().zip(row_vec.iter()).for_each(|(dst, src)| *dst = *src);
28    }
29    
30    // Column-wise FFT
31    let col_fft = planner.plan_fft_forward(rows);
32    for mut col in output.columns_mut() {
33        // Convert column to Vec to ensure contiguous memory
34        let mut col_vec = col.to_vec();
35        col_fft.process(&mut col_vec);
36        // Copy back to column
37        col.iter_mut().zip(col_vec.iter()).for_each(|(dst, src)| *dst = *src);
38    }
39    
40    output
41}
42
43/// Perform 2D inverse FFT on a complex-valued array.
44/// 
45/// This function performs a 2D inverse Fast Fourier Transform using row-wise and column-wise
46/// 1D inverse FFTs. The result is normalized by dividing by the total number of elements.
47/// 
48/// # Arguments
49/// * `input` - Complex-valued 2D array to transform
50/// 
51/// # Returns
52/// A new 2D array containing the inverse Fourier transform
53pub fn ifft2d(input: &ArrayView2<Complex64>) -> Array2<Complex64> {
54    let (rows, cols) = input.dim();
55    let mut output = input.to_owned();
56    let mut planner = FftPlanner::new();
57    
58    // Row-wise IFFT
59    let row_fft = planner.plan_fft_inverse(cols);
60    for mut row in output.rows_mut() {
61        // Convert row to Vec to ensure contiguous memory
62        let mut row_vec = row.to_vec();
63        row_fft.process(&mut row_vec);
64        // Copy back to row
65        row.iter_mut().zip(row_vec.iter()).for_each(|(dst, src)| *dst = *src);
66    }
67    
68    // Column-wise IFFT
69    let col_fft = planner.plan_fft_inverse(rows);
70    for mut col in output.columns_mut() {
71        // Convert column to Vec to ensure contiguous memory
72        let mut col_vec = col.to_vec();
73        col_fft.process(&mut col_vec);
74        // Copy back to column
75        col.iter_mut().zip(col_vec.iter()).for_each(|(dst, src)| *dst = *src);
76    }
77    
78    // Normalize
79    let n = (rows * cols) as f64;
80    output.mapv_inplace(|x| x / n);
81    
82    output
83}
84
85/// Shift zero-frequency component to center of spectrum.
86/// 
87/// This function rearranges the output of an FFT so that the zero-frequency component
88/// is at the center of the spectrum. This is useful for visualization and further processing.
89/// 
90/// # Arguments
91/// * `input` - Complex-valued 2D array to shift
92/// 
93/// # Returns
94/// A new 2D array with the zero-frequency component at the center
95pub fn fftshift(input: &ArrayView2<Complex64>) -> Array2<Complex64> {
96    let (rows, cols) = input.dim();
97    let row_mid = rows / 2;
98    let col_mid = cols / 2;
99    
100    let mut output = Array2::zeros((rows, cols));
101    
102    // Top-left quadrant → Bottom-right
103    output.slice_mut(s![row_mid.., col_mid..])
104        .assign(&input.slice(s![..row_mid, ..col_mid]));
105    
106    // Top-right quadrant → Bottom-left
107    output.slice_mut(s![row_mid.., ..col_mid])
108        .assign(&input.slice(s![..row_mid, col_mid..]));
109    
110    // Bottom-left quadrant → Top-right
111    output.slice_mut(s![..row_mid, col_mid..])
112        .assign(&input.slice(s![row_mid.., ..col_mid]));
113    
114    // Bottom-right quadrant → Top-left
115    output.slice_mut(s![..row_mid, ..col_mid])
116        .assign(&input.slice(s![row_mid.., col_mid..]));
117    
118    output
119}
120
121/// Inverse shift of zero-frequency component from center.
122/// 
123/// This function is the inverse operation of `fftshift`. For even-sized arrays,
124/// this operation is identical to `fftshift`.
125/// 
126/// # Arguments
127/// * `input` - Complex-valued 2D array to shift
128/// 
129/// # Returns
130/// A new 2D array with the zero-frequency component moved from the center
131pub fn ifftshift(input: &ArrayView2<Complex64>) -> Array2<Complex64> {
132    // For even-sized arrays, fftshift and ifftshift are the same operation
133    fftshift(input)
134}
rusty_psf/fft/transform.rs

rusty_psf/fft/
transform.rs
