spectrograms/
fft_backend.rs

1use num_complex::Complex;
2
3use crate::{SpectrogramError, SpectrogramResult};
4
5/// Output size for a real-to-complex FFT of length `n`.
6///
7/// # Arguments
8///
9/// * `n` - Length of the real input
10///
11/// # Returns
12///
13/// Length of the complex output slice required for the FFT.
14#[inline]
15#[must_use]
16pub const fn r2c_output_size(n: usize) -> usize {
17    n / 2 + 1
18}
19/// A planned real-to-complex FFT for a fixed transform length.
20///
21/// Plans must:
22/// - own any internal scratch buffers
23/// - be reusable across many calls
24/// - perform no heap allocation during `process`
25pub trait R2cPlan {
26    fn n_fft(&self) -> usize;
27    fn output_len(&self) -> usize;
28
29    /// Process real input to complex output.
30    ///
31    /// # Arguments
32    ///
33    /// * `input` - Real time domain input
34    /// * `output` - Complex frequency domain output
35    ///
36    /// # Returns
37    ///
38    /// Ok(()) if successful, or SpectrogramError if dimensions are invalid.
39    ///
40    /// # Errors
41    ///
42    /// Returns SpectrogramError::dimension_mismatch if input or output lengths do not match expected sizes.
43    fn process(&mut self, input: &[f64], output: &mut [Complex<f64>]) -> SpectrogramResult<()>;
44}
45
46/// A planned real-to-complex FFT using f32 arithmetic.
47///
48/// Same contract as [`R2cPlan`] but operates on single-precision floats.
49pub trait R2cPlanF32: Send {
50    fn n_fft(&self) -> usize;
51    fn output_len(&self) -> usize;
52    /// Perform the FFT.
53    ///
54    /// # Errors
55    ///
56    /// Returns an error if the input/output lengths don't match the plan size.
57    fn process(&mut self, input: &[f32], output: &mut [Complex<f32>]) -> SpectrogramResult<()>;
58}
59
60/// Planner that can construct FFT plans.
61pub trait R2cPlanner {
62    type Plan: R2cPlan;
63
64    /// Plan a real-to-complex FFT.
65    ///
66    /// # Arguments
67    ///
68    /// * `n_fft` - FFT size
69    ///
70    /// # Returns
71    ///
72    /// A plan that can perform the FFT.
73    ///
74    /// # Errors
75    ///
76    /// Returns SpectrogramError if planning fails.
77    fn plan_r2c(&mut self, n_fft: usize) -> SpectrogramResult<Self::Plan>;
78}
79
80/// A planned complex-to-real inverse FFT for a fixed transform length.
81pub trait C2rPlan {
82    /// FFT size
83    ///
84    /// # Returns
85    ///
86    /// Length of the FFT.
87    fn n_fft(&self) -> usize;
88
89    /// Input length for complex-to-real inverse FFT is r2c_output_size(n_fft)
90    ///
91    /// # Returns
92    ///
93    /// Length of the complex input slice required for the inverse FFT.
94    fn input_len(&self) -> usize;
95
96    /// Process complex input to real output.
97    ///
98    /// # Arguments
99    ///
100    /// * `input` - Complex frequency domain input
101    /// * `output` - Real time domain output
102    ///
103    /// # Returns
104    ///
105    /// Ok(()) if successful, or SpectrogramError if dimensions are invalid.
106    ///
107    /// # Errors
108    ///
109    /// Returns SpectrogramError::dimension_mismatch if input or output lengths do not match expected sizes.
110    fn process(&mut self, input: &[Complex<f64>], output: &mut [f64]) -> SpectrogramResult<()>;
111}
112
113/// Planner that can construct inverse FFT plans.
114pub trait C2rPlanner {
115    type Plan: C2rPlan;
116
117    /// Plan a complex-to-real inverse FFT.
118    ///
119    /// # Arguments
120    ///
121    /// * `n_fft` - FFT size
122    ///
123    /// # Returns
124    ///
125    /// A plan that can perform the inverse FFT.
126    ///
127    /// # Errors
128    ///
129    /// Returns SpectrogramError if planning fails.
130    fn plan_c2r(&mut self, n_fft: usize) -> SpectrogramResult<Self::Plan>;
131}
132
133/// A planned complex-to-complex FFT for a fixed transform length.
134///
135/// The transform is **in-place**: input is overwritten with output.
136/// Plans must own internal scratch buffers and perform no heap allocation during `forward`/`inverse`.
137///
138/// Normalization: neither `forward` nor `inverse` normalizes. The caller is responsible
139/// for dividing by N after an inverse transform if a round-trip is desired.
140pub trait C2cPlan: Send {
141    fn n_fft(&self) -> usize;
142    /// Forward DFT: `X[k] = Σ x[n] exp(−2πi·k·n/N)`
143    ///
144    /// # Errors
145    ///
146    /// Returns an error if the buffer length doesn't match the plan size.
147    fn forward(&mut self, buf: &mut [Complex<f64>]) -> SpectrogramResult<()>;
148    /// Inverse DFT (unnormalized): `x[n] = Σ X[k] exp(+2πi·k·n/N)`
149    ///
150    /// # Errors
151    ///
152    /// Returns an error if the buffer length doesn't match the plan size.
153    fn inverse(&mut self, buf: &mut [Complex<f64>]) -> SpectrogramResult<()>;
154}
155
156/// A planned complex-to-complex FFT using f32 (single-precision) arithmetic.
157///
158/// Same contract as [`C2cPlan`] but operates on single-precision floats.
159pub trait C2cPlanF32: Send {
160    fn n_fft(&self) -> usize;
161    /// Forward DFT (f32).
162    ///
163    /// # Errors
164    ///
165    /// Returns an error if the buffer length doesn't match the plan size.
166    fn forward(&mut self, buf: &mut [Complex<f32>]) -> SpectrogramResult<()>;
167    /// Inverse DFT, unnormalized (f32).
168    ///
169    /// # Errors
170    ///
171    /// Returns an error if the buffer length doesn't match the plan size.
172    fn inverse(&mut self, buf: &mut [Complex<f32>]) -> SpectrogramResult<()>;
173}
174
175// ============================================================================
176// 2D FFT Traits
177// ============================================================================
178
179/// Output size for a 2D real-to-complex FFT of dimensions (nrows, ncols).
180/// The output has shape (nrows, ncols/2 + 1) due to Hermitian symmetry.
181///
182/// # Arguments
183///
184/// * `nrows` - Number of rows in the input
185/// * `ncols` - Number of columns in the input
186///
187/// # Returns
188///
189/// A tuple (out_rows, out_cols) representing the output dimensions.
190#[inline]
191#[must_use]
192pub const fn r2c_output_size_2d(nrows: usize, ncols: usize) -> (usize, usize) {
193    (nrows, ncols / 2 + 1)
194}
195
196/// A planned 2D real-to-complex FFT for fixed dimensions.
197///
198/// Plans must:
199/// - own any internal scratch buffers
200/// - be reusable across many calls
201/// - perform no heap allocation during `process`
202pub trait R2cPlan2d {
203    /// Process 2D real input to complex output.
204    ///
205    /// # Arguments
206    ///
207    /// * `input` - nrows x ncols real values (row-major, flat slice)
208    /// * `output` - nrows x (ncols/2 + 1) complex values (row-major, flat slice)
209    ///
210    /// # Returns
211    ///
212    /// Ok(()) if successful, or SpectrogramError if dimensions are invalid.
213    ///
214    /// # Errors
215    ///
216    /// Returns SpectrogramError::dimension_mismatch if input or output lengths do not match expected sizes.
217    fn process(&mut self, input: &[f64], output: &mut [Complex<f64>]) -> SpectrogramResult<()>;
218}
219
220/// Planner that can construct 2D FFT plans.
221pub trait R2cPlanner2d {
222    type Plan: R2cPlan2d;
223
224    /// Plan a 2D real-to-complex FFT.
225    ///
226    ///  # Arguments
227    ///
228    /// * `nrows` - Number of rows in the input
229    /// * `ncols` - Number of columns in the input
230    ///
231    /// # Returns
232    ///
233    /// A plan that can perform the 2D FFT.
234    ///
235    /// # Errors
236    ///
237    /// Returns SpectrogramError if planning fails.
238    fn plan_r2c_2d(&mut self, nrows: usize, ncols: usize) -> SpectrogramResult<Self::Plan>;
239}
240
241/// A planned 2D complex-to-real inverse FFT for fixed dimensions.
242pub trait C2rPlan2d {
243    /// Process 2D complex input to real output.
244    ///
245    /// # Arguments
246    ///
247    /// * `input` - nrows x (ncols/2 + 1) complex values (row-major, flat slice)
248    /// * `output` - nrows x ncols real values (row-major, flat slice)
249    ///
250    /// # Returns
251    ///
252    /// Ok(()) if successful, or SpectrogramError if dimensions are invalid.
253    ///
254    /// # Errors
255    ///
256    /// Returns SpectrogramError::dimension_mismatch if input or output lengths do not match expected sizes.
257    fn process(&mut self, input: &[Complex<f64>], output: &mut [f64]) -> SpectrogramResult<()>;
258}
259
260/// Planner that can construct 2D inverse FFT plans.
261pub trait C2rPlanner2d {
262    type Plan: C2rPlan2d;
263
264    /// Plan a 2D complex-to-real inverse FFT.
265    ///
266    /// # Arguments
267    ///
268    /// * `nrows` - Number of rows in the output
269    /// * `ncols` - Number of columns in the output
270    ///
271    /// # Returns
272    ///
273    /// A plan that can perform the 2D inverse FFT.
274    ///
275    /// # Errors
276    ///
277    /// Returns SpectrogramError if planning fails.
278    fn plan_c2r_2d(&mut self, nrows: usize, ncols: usize) -> SpectrogramResult<Self::Plan>;
279}
280
281/// Validate 1D FFT input/output dimensions.
282///
283/// # Arguments
284///
285/// * `n_fft` - FFT size
286/// * `input` - Input slice of real values
287/// * `output` - Output slice of complex values
288///
289/// # Returns
290///
291/// Ok(()) if dimensions are valid, or SpectrogramError if not.
292///
293/// # Errors
294///
295/// Returns SpectrogramError::dimension_mismatch if input or output lengths do not match expected sizes.
296#[inline]
297pub const fn validate_fft_io(
298    n_fft: usize,
299    input: &[f64],
300    output: &[Complex<f64>],
301) -> SpectrogramResult<()> {
302    if input.len() != n_fft {
303        return Err(SpectrogramError::dimension_mismatch(n_fft, input.len()));
304    }
305
306    let expected_out = r2c_output_size(n_fft);
307    if output.len() != expected_out {
308        return Err(SpectrogramError::dimension_mismatch(
309            expected_out,
310            output.len(),
311        ));
312    }
313
314    Ok(())
315}
316
317/// Validate 2D FFT input/output dimensions.
318///
319/// # Arguments
320///
321/// * `nrows` - Number of rows in the input
322/// * `ncols` - Number of columns in the input
323/// * `input` - Input slice of real values
324/// * `output` - Output slice of complex values
325///
326/// # Returns
327///
328/// Ok(()) if dimensions are valid, or SpectrogramError if not.
329///
330/// # Errors
331///
332/// Returns SpectrogramError::dimension_mismatch if input or output lengths do not match expected sizes.
333#[inline]
334pub const fn validate_fft_io_2d(
335    nrows: usize,
336    ncols: usize,
337    input: &[f64],
338    output: &[Complex<f64>],
339) -> SpectrogramResult<()> {
340    let input_len = nrows * ncols;
341    if input.len() != input_len {
342        return Err(SpectrogramError::dimension_mismatch(input_len, input.len()));
343    }
344
345    let (out_rows, out_cols) = r2c_output_size_2d(nrows, ncols);
346    let output_len = out_rows * out_cols;
347    if output.len() != output_len {
348        return Err(SpectrogramError::dimension_mismatch(
349            output_len,
350            output.len(),
351        ));
352    }
353
354    Ok(())
355}
356
357#[cfg(feature = "realfft")]
358pub mod realfft_backend {
359    use std::collections::HashMap;
360    use std::sync::Arc;
361
362    use num_complex::Complex;
363    pub use realfft::{ComplexToReal, RealFftPlanner as InnerPlanner, RealToComplex};
364
365    use crate::fft_backend::{
366        C2rPlan, C2rPlanner, R2cPlan, R2cPlanner, r2c_output_size, validate_fft_io,
367    };
368    use crate::{SpectrogramError, SpectrogramResult};
369
370    /// RealFftPlanner
371    ///
372    /// Wraps the realfft::RealFftPlanner and adds caching for created plans.
373    ///
374    /// This planner maintains separate caches for forward (real-to-complex)
375    /// and inverse (complex-to-real) FFT plans to avoid redundant plan creation.
376    #[derive(Default)]
377    pub struct RealFftPlanner {
378        inner: InnerPlanner<f64>,
379        cache_r2c: HashMap<usize, Arc<dyn RealToComplex<f64>>>,
380        cache_c2r: HashMap<usize, Arc<dyn ComplexToReal<f64>>>,
381    }
382
383    impl RealFftPlanner {
384        /// Create a new RealFftPlanner with empty caches.
385        ///
386        /// # Returns
387        ///
388        /// A new instance of `RealFftPlanner` with empty caches.
389        #[inline]
390        #[must_use]
391        pub fn new() -> Self {
392            Self::default()
393        }
394
395        pub(crate) fn get_or_create(&mut self, n_fft: usize) -> Arc<dyn RealToComplex<f64>> {
396            if let Some(p) = self.cache_r2c.get(&n_fft) {
397                return p.clone();
398            }
399            let p = self.inner.plan_fft_forward(n_fft);
400            self.cache_r2c.insert(n_fft, p.clone());
401            p
402        }
403
404        pub(crate) fn get_or_create_inverse(
405            &mut self,
406            n_fft: usize,
407        ) -> Arc<dyn ComplexToReal<f64>> {
408            if let Some(p) = self.cache_c2r.get(&n_fft) {
409                return p.clone();
410            }
411            let p = self.inner.plan_fft_inverse(n_fft);
412            self.cache_c2r.insert(n_fft, p.clone());
413            p
414        }
415    }
416
417    /// Real-to-Complex FFT Plan
418    ///
419    /// Implements a plan for performing FFTs from real time domain to complex frequency domain.=
420    #[derive(Clone)]
421    pub struct RealFftPlan {
422        n_fft: usize,
423        plan: Arc<dyn RealToComplex<f64>>,
424        scratch: Vec<f64>,
425    }
426
427    impl RealFftPlan {
428        pub(crate) fn new(n_fft: usize, plan: Arc<dyn RealToComplex<f64>>) -> Self {
429            Self {
430                n_fft,
431                plan,
432                scratch: vec![0.0; n_fft],
433            }
434        }
435    }
436
437    impl R2cPlan for RealFftPlan {
438        #[inline]
439        fn n_fft(&self) -> usize {
440            self.n_fft
441        }
442
443        #[inline]
444        fn output_len(&self) -> usize {
445            r2c_output_size(self.n_fft)
446        }
447
448        #[inline]
449        fn process(&mut self, input: &[f64], output: &mut [Complex<f64>]) -> SpectrogramResult<()> {
450            validate_fft_io(self.n_fft, input, output)?;
451
452            self.scratch.copy_from_slice(input);
453
454            self.plan
455                .process(&mut self.scratch, output)
456                .map_err(|e| SpectrogramError::fft_backend("realfft", format!("{e:?}")))
457        }
458    }
459
460    impl R2cPlanner for RealFftPlanner {
461        type Plan = RealFftPlan;
462
463        #[inline]
464        fn plan_r2c(&mut self, n_fft: usize) -> SpectrogramResult<Self::Plan> {
465            let plan = self.get_or_create(n_fft);
466            Ok(RealFftPlan::new(n_fft, plan))
467        }
468    }
469
470    /// Real-to-Complex FFT plan using f32 (single-precision) arithmetic.
471    pub struct RealFftPlanF32 {
472        n_fft: usize,
473        plan: Arc<dyn RealToComplex<f32>>,
474        scratch: Vec<f32>,
475    }
476
477    impl RealFftPlanF32 {
478        /// Create an f32 R2c plan for transform size `n_fft`.
479        #[must_use]
480        pub fn new(n_fft: usize) -> Self {
481            let mut planner: InnerPlanner<f32> = InnerPlanner::new();
482            let plan = planner.plan_fft_forward(n_fft);
483            Self {
484                n_fft,
485                plan,
486                scratch: vec![0.0f32; n_fft],
487            }
488        }
489    }
490
491    impl crate::fft_backend::R2cPlanF32 for RealFftPlanF32 {
492        #[inline]
493        fn n_fft(&self) -> usize {
494            self.n_fft
495        }
496
497        #[inline]
498        fn output_len(&self) -> usize {
499            r2c_output_size(self.n_fft)
500        }
501
502        #[inline]
503        fn process(&mut self, input: &[f32], output: &mut [Complex<f32>]) -> SpectrogramResult<()> {
504            if input.len() != self.n_fft {
505                return Err(SpectrogramError::dimension_mismatch(
506                    self.n_fft,
507                    input.len(),
508                ));
509            }
510            self.scratch.copy_from_slice(input);
511            self.plan
512                .process(&mut self.scratch, output)
513                .map_err(|e| SpectrogramError::fft_backend("realfft_f32", format!("{e:?}")))
514        }
515    }
516
517    // ── Complex-to-Complex plans (f64 and f32) ────────────────────────────────
518
519    /// Complex-to-complex FFT plan (f64) backed by rustfft.
520    ///
521    /// Both forward and inverse directions share a single scratch buffer sized to
522    /// the larger of the two inplace scratch requirements.
523    pub struct RealFftC2cPlan {
524        n_fft: usize,
525        fwd: Arc<dyn rustfft::Fft<f64>>,
526        inv: Arc<dyn rustfft::Fft<f64>>,
527        scratch: Vec<Complex<f64>>,
528    }
529
530    impl RealFftC2cPlan {
531        /// Create a C2c plan for transform size `n_fft`.
532        #[must_use]
533        pub fn new(n_fft: usize) -> Self {
534            let mut p = RustFftPlanner::<f64>::new();
535            let fwd = p.plan_fft_forward(n_fft);
536            let inv = p.plan_fft_inverse(n_fft);
537            let scratch_len = fwd
538                .get_inplace_scratch_len()
539                .max(inv.get_inplace_scratch_len());
540            Self {
541                n_fft,
542                fwd,
543                inv,
544                scratch: vec![Complex::new(0.0, 0.0); scratch_len],
545            }
546        }
547    }
548
549    impl crate::fft_backend::C2cPlan for RealFftC2cPlan {
550        #[inline]
551        fn n_fft(&self) -> usize {
552            self.n_fft
553        }
554
555        #[inline]
556        fn forward(&mut self, buf: &mut [Complex<f64>]) -> SpectrogramResult<()> {
557            if buf.len() != self.n_fft {
558                return Err(SpectrogramError::dimension_mismatch(self.n_fft, buf.len()));
559            }
560            self.fwd.process_with_scratch(buf, &mut self.scratch);
561            Ok(())
562        }
563
564        #[inline]
565        fn inverse(&mut self, buf: &mut [Complex<f64>]) -> SpectrogramResult<()> {
566            if buf.len() != self.n_fft {
567                return Err(SpectrogramError::dimension_mismatch(self.n_fft, buf.len()));
568            }
569            self.inv.process_with_scratch(buf, &mut self.scratch);
570            Ok(())
571        }
572    }
573
574    // SAFETY: rustfft::Fft<T> is Send; scratch is owned by the plan.
575    unsafe impl Send for RealFftC2cPlan {}
576
577    /// Complex-to-complex FFT plan (f32) backed by rustfft.
578    pub struct RealFftC2cPlanF32 {
579        n_fft: usize,
580        fwd: Arc<dyn rustfft::Fft<f32>>,
581        inv: Arc<dyn rustfft::Fft<f32>>,
582        scratch: Vec<Complex<f32>>,
583    }
584
585    impl RealFftC2cPlanF32 {
586        /// Create an f32 C2c plan for transform size `n_fft`.
587        #[must_use]
588        pub fn new(n_fft: usize) -> Self {
589            let mut p = RustFftPlanner::<f32>::new();
590            let fwd = p.plan_fft_forward(n_fft);
591            let inv = p.plan_fft_inverse(n_fft);
592            let scratch_len = fwd
593                .get_inplace_scratch_len()
594                .max(inv.get_inplace_scratch_len());
595            Self {
596                n_fft,
597                fwd,
598                inv,
599                scratch: vec![Complex::new(0.0f32, 0.0f32); scratch_len],
600            }
601        }
602    }
603
604    impl crate::fft_backend::C2cPlanF32 for RealFftC2cPlanF32 {
605        #[inline]
606        fn n_fft(&self) -> usize {
607            self.n_fft
608        }
609
610        #[inline]
611        fn forward(&mut self, buf: &mut [Complex<f32>]) -> SpectrogramResult<()> {
612            if buf.len() != self.n_fft {
613                return Err(SpectrogramError::dimension_mismatch(self.n_fft, buf.len()));
614            }
615            self.fwd.process_with_scratch(buf, &mut self.scratch);
616            Ok(())
617        }
618
619        #[inline]
620        fn inverse(&mut self, buf: &mut [Complex<f32>]) -> SpectrogramResult<()> {
621            if buf.len() != self.n_fft {
622                return Err(SpectrogramError::dimension_mismatch(self.n_fft, buf.len()));
623            }
624            self.inv.process_with_scratch(buf, &mut self.scratch);
625            Ok(())
626        }
627    }
628
629    // SAFETY: rustfft::Fft<T> is Send; scratch is owned by the plan.
630    unsafe impl Send for RealFftC2cPlanF32 {}
631
632    /// Complex-to-Real Inverse FFT Plan
633    ///
634    /// Implements a plan for performing inverse FFTs from complex frequency domain
635    /// to real time domain.
636    #[derive(Clone)]
637    pub struct RealFftInversePlan {
638        n_fft: usize,
639        plan: Arc<dyn ComplexToReal<f64>>,
640        scratch: Vec<Complex<f64>>,
641    }
642
643    impl RealFftInversePlan {
644        pub(crate) fn new(n_fft: usize, plan: Arc<dyn ComplexToReal<f64>>) -> Self {
645            let scratch_len = r2c_output_size(n_fft);
646            Self {
647                n_fft,
648                plan,
649                scratch: vec![Complex::new(0.0, 0.0); scratch_len],
650            }
651        }
652    }
653
654    impl C2rPlan for RealFftInversePlan {
655        #[inline]
656        fn n_fft(&self) -> usize {
657            self.n_fft
658        }
659
660        #[inline]
661        fn input_len(&self) -> usize {
662            r2c_output_size(self.n_fft)
663        }
664
665        #[inline]
666        fn process(&mut self, input: &[Complex<f64>], output: &mut [f64]) -> SpectrogramResult<()> {
667            let expected_in = r2c_output_size(self.n_fft);
668            if input.len() != expected_in {
669                return Err(SpectrogramError::dimension_mismatch(
670                    expected_in,
671                    input.len(),
672                ));
673            }
674            if output.len() != self.n_fft {
675                return Err(SpectrogramError::dimension_mismatch(
676                    self.n_fft,
677                    output.len(),
678                ));
679            }
680
681            self.scratch.copy_from_slice(input);
682
683            self.plan
684                .process(&mut self.scratch, output)
685                .map_err(|e| SpectrogramError::fft_backend("realfft", format!("{e:?}")))?;
686
687            // RealFFT inverse needs normalization
688            let scale = 1.0 / self.n_fft as f64;
689            for val in output.iter_mut() {
690                *val *= scale;
691            }
692
693            Ok(())
694        }
695    }
696
697    impl C2rPlanner for RealFftPlanner {
698        type Plan = RealFftInversePlan;
699
700        #[inline]
701        fn plan_c2r(&mut self, n_fft: usize) -> SpectrogramResult<Self::Plan> {
702            let plan = self.get_or_create_inverse(n_fft);
703            Ok(RealFftInversePlan::new(n_fft, plan))
704        }
705    }
706
707    // ========================================================================
708    // 2D FFT Implementation (Row-Column Decomposition)
709    // ========================================================================
710
711    #[cfg(feature = "realfft")]
712    use rustfft::FftPlanner as RustFftPlanner;
713
714    impl RealFftPlanner {
715        /// Get or create a complex-to-complex FFT plan for column transforms.
716        pub(crate) fn get_or_create_complex(n: usize) -> Arc<dyn rustfft::Fft<f64>> {
717            // Use the inner RealFftPlanner's complex FFT planner
718            // Note: realfft::RealFftPlanner wraps rustfft::FftPlanner
719            // We need to create a separate planner for complex FFTs
720            let mut complex_planner = RustFftPlanner::<f64>::new();
721            complex_planner.plan_fft_forward(n)
722        }
723
724        pub(crate) fn get_or_create_complex_inverse(n: usize) -> Arc<dyn rustfft::Fft<f64>> {
725            let mut complex_planner = RustFftPlanner::<f64>::new();
726            complex_planner.plan_fft_inverse(n)
727        }
728    }
729
730    /// 2D Real-to-Complex FFT Plan
731    ///
732    /// Implements row-column decomposition:
733    ///
734    /// 1. Perform 1D real-to-complex FFTs on each row
735    /// 2. Perform 1D complex-to-complex FFTs on each column of the intermediate result
736    ///
737    /// This plan owns scratch buffers for row and column transforms, as well as
738    /// an intermediate buffer to hold the results after the row transforms.
739    ///
740    /// The process method performs the full 2D FFT in-place without additional allocations.
741    #[derive(Clone)]
742    pub struct RealFftPlan2d {
743        nrows: usize,
744        ncols: usize,
745        out_shape: (usize, usize),
746        // Plans for row transforms (real -> complex, size ncols)
747        row_plan: Arc<dyn RealToComplex<f64>>,
748        // Plans for column transforms (complex -> complex, size nrows)
749        col_plan: Arc<dyn rustfft::Fft<f64>>,
750        // Scratch buffers
751        scratch_row: Vec<f64>,
752        scratch_col: Vec<Complex<f64>>,
753        // Intermediate storage after row FFTs
754        intermediate: Vec<Complex<f64>>,
755    }
756
757    impl RealFftPlan2d {
758        pub(crate) fn new(
759            nrows: usize,
760            ncols: usize,
761            row_plan: Arc<dyn RealToComplex<f64>>,
762            col_plan: Arc<dyn rustfft::Fft<f64>>,
763        ) -> Self {
764            let out_shape = crate::fft_backend::r2c_output_size_2d(nrows, ncols);
765            let intermediate_len = nrows * out_shape.1;
766
767            Self {
768                nrows,
769                ncols,
770                out_shape,
771                row_plan,
772                col_plan,
773                scratch_row: vec![0.0; ncols],
774                scratch_col: vec![Complex::new(0.0, 0.0); nrows],
775                intermediate: vec![Complex::new(0.0, 0.0); intermediate_len],
776            }
777        }
778    }
779
780    impl crate::fft_backend::R2cPlan2d for RealFftPlan2d {
781        fn process(&mut self, input: &[f64], output: &mut [Complex<f64>]) -> SpectrogramResult<()> {
782            crate::fft_backend::validate_fft_io_2d(self.nrows, self.ncols, input, output)?;
783
784            // Step 1: FFT each row (real -> complex)
785            for row_idx in 0..self.nrows {
786                let row_start = row_idx * self.ncols;
787                let row_end = row_start + self.ncols;
788                self.scratch_row.copy_from_slice(&input[row_start..row_end]);
789
790                let out_start = row_idx * self.out_shape.1;
791                let out_end = out_start + self.out_shape.1;
792
793                self.row_plan
794                    .process(
795                        &mut self.scratch_row,
796                        &mut self.intermediate[out_start..out_end],
797                    )
798                    .map_err(|e| SpectrogramError::fft_backend("realfft", format!("{e:?}")))?;
799            }
800
801            // Step 2: FFT each column (complex -> complex)
802            for col_idx in 0..self.out_shape.1 {
803                // Extract column from intermediate buffer
804                for row_idx in 0..self.nrows {
805                    self.scratch_col[row_idx] =
806                        self.intermediate[row_idx * self.out_shape.1 + col_idx];
807                }
808
809                // FFT the column
810                self.col_plan.process(&mut self.scratch_col);
811
812                // Write column back to output
813                for row_idx in 0..self.nrows {
814                    output[row_idx * self.out_shape.1 + col_idx] = self.scratch_col[row_idx];
815                }
816            }
817
818            Ok(())
819        }
820    }
821
822    impl crate::fft_backend::R2cPlanner2d for RealFftPlanner {
823        type Plan = RealFftPlan2d;
824
825        fn plan_r2c_2d(&mut self, nrows: usize, ncols: usize) -> SpectrogramResult<Self::Plan> {
826            let row_plan = self.get_or_create(ncols);
827            let col_plan = Self::get_or_create_complex(nrows);
828            Ok(RealFftPlan2d::new(nrows, ncols, row_plan, col_plan))
829        }
830    }
831
832    #[derive(Clone)]
833    pub struct RealFftInversePlan2d {
834        nrows: usize,
835        ncols: usize,
836        in_shape: (usize, usize),
837        // Plans for column inverse transforms (complex -> complex, size nrows)
838        col_plan: Arc<dyn rustfft::Fft<f64>>,
839        // Plans for row inverse transforms (complex -> real, size ncols)
840        row_plan: Arc<dyn ComplexToReal<f64>>,
841        // Scratch buffers
842        scratch_col: Vec<Complex<f64>>,
843        scratch_row: Vec<Complex<f64>>,
844        // Intermediate storage after column IFFTs
845        intermediate: Vec<Complex<f64>>,
846    }
847
848    impl RealFftInversePlan2d {
849        pub(crate) fn new(
850            nrows: usize,
851            ncols: usize,
852            col_plan: Arc<dyn rustfft::Fft<f64>>,
853            row_plan: Arc<dyn ComplexToReal<f64>>,
854        ) -> Self {
855            let in_shape = crate::fft_backend::r2c_output_size_2d(nrows, ncols);
856            let intermediate_len = nrows * in_shape.1;
857
858            Self {
859                nrows,
860                ncols,
861                in_shape,
862                col_plan,
863                row_plan,
864                scratch_col: vec![Complex::new(0.0, 0.0); nrows],
865                scratch_row: vec![Complex::new(0.0, 0.0); in_shape.1],
866                intermediate: vec![Complex::new(0.0, 0.0); intermediate_len],
867            }
868        }
869    }
870
871    impl crate::fft_backend::C2rPlan2d for RealFftInversePlan2d {
872        fn process(&mut self, input: &[Complex<f64>], output: &mut [f64]) -> SpectrogramResult<()> {
873            let expected_in_len = self.in_shape.0 * self.in_shape.1;
874            if input.len() != expected_in_len {
875                return Err(SpectrogramError::dimension_mismatch(
876                    expected_in_len,
877                    input.len(),
878                ));
879            }
880
881            let expected_out_len = self.nrows * self.ncols;
882            if output.len() != expected_out_len {
883                return Err(SpectrogramError::dimension_mismatch(
884                    expected_out_len,
885                    output.len(),
886                ));
887            }
888
889            // Copy input to intermediate buffer
890            self.intermediate.copy_from_slice(input);
891
892            // Step 1: Inverse FFT each column (complex -> complex)
893            for col_idx in 0..self.in_shape.1 {
894                // Extract column
895                for row_idx in 0..self.nrows {
896                    self.scratch_col[row_idx] =
897                        self.intermediate[row_idx * self.in_shape.1 + col_idx];
898                }
899
900                // IFFT the column
901                self.col_plan.process(&mut self.scratch_col);
902
903                // Write column back
904                for row_idx in 0..self.nrows {
905                    self.intermediate[row_idx * self.in_shape.1 + col_idx] =
906                        self.scratch_col[row_idx];
907                }
908            }
909
910            // Enforce Hermitian symmetry: DC (column 0) and Nyquist (last column if even)
911            // must have purely real values for each row
912            for row_idx in 0..self.nrows {
913                // DC component (first column) must be real
914                self.intermediate[row_idx * self.in_shape.1].im = 0.0;
915
916                // Nyquist component (last column, if ncols is even) must be real
917                if self.ncols.is_multiple_of(2) {
918                    let nyquist_col = self.in_shape.1 - 1;
919                    self.intermediate[row_idx * self.in_shape.1 + nyquist_col].im = 0.0;
920                }
921            }
922
923            // Step 2: Inverse FFT each row (complex -> real)
924            for row_idx in 0..self.nrows {
925                let row_start = row_idx * self.in_shape.1;
926                let row_end = row_start + self.in_shape.1;
927                self.scratch_row
928                    .copy_from_slice(&self.intermediate[row_start..row_end]);
929
930                let out_start = row_idx * self.ncols;
931                let out_end = out_start + self.ncols;
932
933                self.row_plan
934                    .process(&mut self.scratch_row, &mut output[out_start..out_end])
935                    .map_err(|e| SpectrogramError::fft_backend("realfft", format!("{e:?}")))?;
936            }
937
938            // RealFFT inverse needs normalization
939            let scale = 1.0 / (self.nrows * self.ncols) as f64;
940            for val in output.iter_mut() {
941                *val *= scale;
942            }
943
944            Ok(())
945        }
946    }
947
948    impl crate::fft_backend::C2rPlanner2d for RealFftPlanner {
949        type Plan = RealFftInversePlan2d;
950
951        fn plan_c2r_2d(&mut self, nrows: usize, ncols: usize) -> SpectrogramResult<Self::Plan> {
952            let col_plan = Self::get_or_create_complex_inverse(nrows);
953            let row_plan = self.get_or_create_inverse(ncols);
954            Ok(RealFftInversePlan2d::new(nrows, ncols, col_plan, row_plan))
955        }
956    }
957}
958
959// ============================================================================
960// Global Plan Cache (for one-shot FFT functions)
961// ============================================================================
962
963#[cfg(feature = "realfft")]
964pub mod plan_cache {
965    use super::{C2rPlanner, R2cPlanner, SpectrogramError, SpectrogramResult};
966
967    use std::collections::HashMap;
968    use std::sync::{Arc, Mutex};
969
970    use crate::fft_backend::realfft_backend::{RealFftInversePlan, RealFftPlan, RealFftPlanner};
971
972    /// Global cache of forward FFT plans, keyed by n_fft size.
973    /// Uses Arc for cheap cloning and Mutex for thread safety.
974    static R2C_PLAN_CACHE: std::sync::LazyLock<Mutex<HashMap<usize, Arc<RealFftPlan>>>> =
975        std::sync::LazyLock::new(|| Mutex::new(HashMap::new()));
976
977    /// Global cache of inverse FFT plans, keyed by n_fft size.
978    static C2R_PLAN_CACHE: std::sync::LazyLock<Mutex<HashMap<usize, Arc<RealFftInversePlan>>>> =
979        std::sync::LazyLock::new(|| Mutex::new(HashMap::new()));
980
981    /// Maximum number of cached plans to prevent unbounded memory growth.
982    /// Plans are evicted FIFO when this limit is reached.
983    const MAX_CACHED_PLANS: usize = 100;
984
985    /// Get or create a forward (real-to-complex) FFT plan from the global cache.
986    ///
987    /// This function is thread-safe and will return a cached plan if one exists
988    /// for the given size, or create a new one and cache it for future use.
989    ///
990    /// # Arguments
991    ///
992    /// * `n_fft` - FFT size
993    ///
994    /// # Returns
995    ///
996    /// An Arc-wrapped plan that can be cloned cheaply for use across threads.
997    pub fn get_or_create_r2c_plan(n_fft: usize) -> SpectrogramResult<Arc<RealFftPlan>> {
998        let mut cache = R2C_PLAN_CACHE.lock().map_err(|e| {
999            SpectrogramError::fft_backend("plan_cache", format!("mutex poisoned: {e}"))
1000        })?;
1001
1002        // Fast path: plan already exists
1003        if let Some(plan) = cache.get(&n_fft) {
1004            return Ok(Arc::clone(plan));
1005        }
1006
1007        // Slow path: create new plan
1008        let mut planner = RealFftPlanner::new();
1009        let plan = planner.plan_r2c(n_fft)?;
1010        let plan = Arc::new(plan);
1011
1012        // Evict an arbitrary plan if cache is full (HashMap iteration order)
1013        if cache.len() >= MAX_CACHED_PLANS {
1014            if let Some(key) = cache.keys().next().copied() {
1015                cache.remove(&key);
1016            }
1017        }
1018
1019        cache.insert(n_fft, Arc::clone(&plan));
1020        drop(cache);
1021        Ok(plan)
1022    }
1023
1024    /// Get or create an inverse (complex-to-real) FFT plan from the global cache.
1025    ///
1026    /// This function is thread-safe and will return a cached plan if one exists
1027    /// for the given size, or create a new one and cache it for future use.
1028    ///
1029    /// # Arguments
1030    ///
1031    /// * `n_fft` - FFT size
1032    ///
1033    /// # Returns
1034    ///
1035    /// An Arc-wrapped plan that can be cloned cheaply for use across threads.
1036    #[inline]
1037    pub fn get_or_create_c2r_plan(n_fft: usize) -> SpectrogramResult<Arc<RealFftInversePlan>> {
1038        let mut cache = C2R_PLAN_CACHE.lock().map_err(|e| {
1039            SpectrogramError::fft_backend("plan_cache", format!("mutex poisoned: {e}"))
1040        })?;
1041
1042        // Fast path: plan already exists
1043        if let Some(plan) = cache.get(&n_fft) {
1044            return Ok(Arc::clone(plan));
1045        }
1046
1047        // Slow path: create new plan
1048        let mut planner = RealFftPlanner::new();
1049        let plan = planner.plan_c2r(n_fft)?;
1050        let plan = Arc::new(plan);
1051
1052        // Evict an arbitrary plan if cache is full (HashMap iteration order)
1053        if cache.len() >= MAX_CACHED_PLANS {
1054            if let Some(key) = cache.keys().next().copied() {
1055                cache.remove(&key);
1056            }
1057        }
1058
1059        cache.insert(n_fft, Arc::clone(&plan));
1060        drop(cache); // Release lock before returning
1061        Ok(plan)
1062    }
1063
1064    /// Clear all cached FFT plans.
1065    ///
1066    /// This is useful for:
1067    /// - Testing and benchmarking
1068    /// - Memory management in long-running applications
1069    /// - Forcing plan recreation with different settings
1070    ///
1071    /// Plans will be automatically recreated on next use.
1072    #[allow(unused)]
1073    #[inline]
1074    pub fn clear_plan_cache() {
1075        if let Ok(mut cache) = R2C_PLAN_CACHE.lock() {
1076            cache.clear();
1077        }
1078        if let Ok(mut cache) = C2R_PLAN_CACHE.lock() {
1079            cache.clear();
1080        }
1081    }
1082
1083    /// Get cache statistics (for monitoring/debugging).
1084    ///
1085    /// Returns (forward_plans_cached, inverse_plans_cached).
1086    #[allow(unused)]
1087    #[inline]
1088    pub fn cache_stats() -> (usize, usize) {
1089        let r2c_count = R2C_PLAN_CACHE.lock().map(|c| c.len()).unwrap_or(0);
1090        let c2r_count = C2R_PLAN_CACHE.lock().map(|c| c.len()).unwrap_or(0);
1091        (r2c_count, c2r_count)
1092    }
1093}
1094
1095#[cfg(feature = "realfft")]
1096pub use plan_cache::{get_or_create_c2r_plan, get_or_create_r2c_plan};
1097
1098#[cfg(all(feature = "realfft", feature = "python"))]
1099pub use plan_cache::{cache_stats, clear_plan_cache};
1100
1101#[cfg(feature = "fftw")]
1102pub mod fftw_backend {
1103    use std::collections::HashMap;
1104    use std::ptr::NonNull;
1105    use std::sync::{Arc, Mutex};
1106
1107    use num_complex::Complex;
1108
1109    use crate::fft_backend::{
1110        C2rPlan, C2rPlanner, R2cPlan, R2cPlanner, r2c_output_size, validate_fft_io,
1111    };
1112    use crate::{SpectrogramError, SpectrogramResult};
1113
1114    // FFTW plan creation is not thread-safe, so we use a global mutex
1115    static FFTW_PLANNER_LOCK: Mutex<()> = Mutex::new(());
1116
1117    #[derive(Debug)]
1118    struct FftwBuffer<T> {
1119        ptr: NonNull<T>,
1120        _len: usize,
1121    }
1122
1123    impl<T> FftwBuffer<T> {
1124        fn allocate(len: usize) -> SpectrogramResult<Self> {
1125            if len == 0 {
1126                return Err(SpectrogramError::invalid_input("buffer length must be > 0"));
1127            }
1128
1129            let bytes = core::mem::size_of::<T>() * len;
1130            let raw = unsafe { fftw_sys::fftw_malloc(bytes) }.cast::<T>();
1131
1132            let ptr = NonNull::new(raw).ok_or_else(|| {
1133                SpectrogramError::fft_backend("fftw", "fftw_malloc returned null")
1134            })?;
1135
1136            Ok(Self { ptr, _len: len })
1137        }
1138
1139        #[inline]
1140        const fn as_ptr(&self) -> *mut T {
1141            self.ptr.as_ptr()
1142        }
1143    }
1144
1145    impl<T> Drop for FftwBuffer<T> {
1146        fn drop(&mut self) {
1147            unsafe {
1148                fftw_sys::fftw_free(self.ptr.as_ptr().cast::<core::ffi::c_void>());
1149            }
1150        }
1151    }
1152
1153    #[derive(Debug)]
1154    pub(crate) struct PlanInner {
1155        n_fft: usize,
1156        out_len: usize,
1157        plan: fftw_sys::fftw_plan,
1158        input: Arc<FftwBuffer<f64>>,
1159        output: Arc<FftwBuffer<fftw_sys::fftw_complex>>,
1160    }
1161
1162    impl Drop for PlanInner {
1163        fn drop(&mut self) {
1164            unsafe {
1165                fftw_sys::fftw_destroy_plan(self.plan);
1166            }
1167        }
1168    }
1169
1170    #[derive(Debug)]
1171    pub(crate) struct InversePlanInner {
1172        n_fft: usize,
1173        in_len: usize,
1174        plan: fftw_sys::fftw_plan,
1175        input: Arc<FftwBuffer<fftw_sys::fftw_complex>>,
1176        output: Arc<FftwBuffer<f64>>,
1177    }
1178
1179    impl Drop for InversePlanInner {
1180        fn drop(&mut self) {
1181            unsafe {
1182                fftw_sys::fftw_destroy_plan(self.plan);
1183            }
1184        }
1185    }
1186
1187    #[derive(Default)]
1188    pub struct FftwPlanner {
1189        cache_r2c: HashMap<usize, Arc<PlanInner>>,
1190        cache_c2r: HashMap<usize, Arc<InversePlanInner>>,
1191    }
1192
1193    impl FftwPlanner {
1194        #[must_use]
1195        pub fn new() -> Self {
1196            Self::default()
1197        }
1198
1199        pub(crate) fn build_plan(n_fft: usize) -> SpectrogramResult<PlanInner> {
1200            let out_len = r2c_output_size(n_fft);
1201
1202            let input = Arc::new(FftwBuffer::<f64>::allocate(n_fft)?);
1203            let output = Arc::new(FftwBuffer::<fftw_sys::fftw_complex>::allocate(out_len)?);
1204
1205            let n_i32: i32 = n_fft
1206                .try_into()
1207                .map_err(|_| SpectrogramError::invalid_input("n_fft too large for FFTW"))?;
1208
1209            // FFTW plan creation is not thread-safe - must be serialized
1210            let _lock = FFTW_PLANNER_LOCK.lock().map_err(|e| {
1211                SpectrogramError::fft_backend("fftw", format!("FFT planner mutex poisoned: {}", e))
1212            })?;
1213
1214            let plan = unsafe {
1215                fftw_sys::fftw_plan_dft_r2c_1d(
1216                    n_i32,
1217                    input.as_ptr(),
1218                    output.as_ptr(),
1219                    fftw_sys::FFTW_ESTIMATE,
1220                )
1221            };
1222
1223            if plan.is_null() {
1224                return Err(SpectrogramError::fft_backend(
1225                    "fftw",
1226                    "failed to create FFTW plan",
1227                ));
1228            }
1229
1230            Ok(PlanInner {
1231                n_fft,
1232                out_len,
1233                plan,
1234                input,
1235                output,
1236            })
1237        }
1238
1239        pub(crate) fn get_or_create(&mut self, n_fft: usize) -> SpectrogramResult<Arc<PlanInner>> {
1240            if let Some(p) = self.cache_r2c.get(&n_fft) {
1241                return Ok(p.clone());
1242            }
1243
1244            let p = Arc::new(Self::build_plan(n_fft)?);
1245            self.cache_r2c.insert(n_fft, p.clone());
1246            Ok(p)
1247        }
1248
1249        pub(crate) fn build_inverse_plan(n_fft: usize) -> SpectrogramResult<InversePlanInner> {
1250            let in_len = r2c_output_size(n_fft);
1251
1252            let input = Arc::new(FftwBuffer::<fftw_sys::fftw_complex>::allocate(in_len)?);
1253            let output = Arc::new(FftwBuffer::<f64>::allocate(n_fft)?);
1254
1255            let n_i32: i32 = n_fft
1256                .try_into()
1257                .map_err(|_| SpectrogramError::invalid_input("n_fft too large for FFTW"))?;
1258
1259            // FFTW plan creation is not thread-safe - must be serialized
1260            let _lock = FFTW_PLANNER_LOCK.lock().map_err(|e| {
1261                SpectrogramError::fft_backend("fftw", format!("FFT planner mutex poisoned: {}", e))
1262            })?;
1263
1264            let plan = unsafe {
1265                fftw_sys::fftw_plan_dft_c2r_1d(
1266                    n_i32,
1267                    input.as_ptr(),
1268                    output.as_ptr(),
1269                    fftw_sys::FFTW_ESTIMATE,
1270                )
1271            };
1272
1273            if plan.is_null() {
1274                return Err(SpectrogramError::fft_backend(
1275                    "fftw",
1276                    "failed to create FFTW inverse plan",
1277                ));
1278            }
1279
1280            Ok(InversePlanInner {
1281                n_fft,
1282                in_len,
1283                plan,
1284                input,
1285                output,
1286            })
1287        }
1288
1289        pub(crate) fn get_or_create_inverse(
1290            &mut self,
1291            n_fft: usize,
1292        ) -> SpectrogramResult<Arc<InversePlanInner>> {
1293            if let Some(p) = self.cache_c2r.get(&n_fft) {
1294                return Ok(p.clone());
1295            }
1296
1297            let p = Arc::new(Self::build_inverse_plan(n_fft)?);
1298            self.cache_c2r.insert(n_fft, p.clone());
1299            Ok(p)
1300        }
1301    }
1302
1303    #[derive(Debug, Clone)]
1304    pub struct FftwPlan {
1305        inner: Arc<PlanInner>,
1306    }
1307
1308    impl FftwPlan {
1309        pub(crate) const fn new(plan: Arc<PlanInner>) -> Self {
1310            Self { inner: plan }
1311        }
1312    }
1313
1314    impl R2cPlan for FftwPlan {
1315        fn n_fft(&self) -> usize {
1316            self.inner.n_fft
1317        }
1318
1319        fn output_len(&self) -> usize {
1320            self.inner.out_len
1321        }
1322
1323        fn process(&mut self, input: &[f64], output: &mut [Complex<f64>]) -> SpectrogramResult<()> {
1324            validate_fft_io(self.inner.n_fft, input, output)?;
1325
1326            unsafe {
1327                core::ptr::copy_nonoverlapping(
1328                    input.as_ptr(),
1329                    self.inner.input.as_ptr(),
1330                    self.inner.n_fft,
1331                );
1332
1333                fftw_sys::fftw_execute_dft_r2c(
1334                    self.inner.plan,
1335                    self.inner.input.as_ptr(),
1336                    self.inner.output.as_ptr(),
1337                );
1338
1339                for i in 0..self.inner.out_len {
1340                    let c = self.inner.output.as_ptr().add(i) as *const f64;
1341                    let re = *c.add(0);
1342                    let im = *c.add(1);
1343                    output[i] = Complex::new(re, im);
1344                }
1345            }
1346
1347            Ok(())
1348        }
1349    }
1350
1351    impl R2cPlanner for FftwPlanner {
1352        type Plan = FftwPlan;
1353
1354        fn plan_r2c(&mut self, n_fft: usize) -> SpectrogramResult<Self::Plan> {
1355            Ok(FftwPlan {
1356                inner: self.get_or_create(n_fft)?,
1357            })
1358        }
1359    }
1360
1361    #[derive(Debug, Clone)]
1362    pub struct FftwInversePlan {
1363        inner: Arc<InversePlanInner>,
1364    }
1365
1366    impl C2rPlan for FftwInversePlan {
1367        fn n_fft(&self) -> usize {
1368            self.inner.n_fft
1369        }
1370
1371        fn input_len(&self) -> usize {
1372            self.inner.in_len
1373        }
1374
1375        fn process(&mut self, input: &[Complex<f64>], output: &mut [f64]) -> SpectrogramResult<()> {
1376            if input.len() != self.inner.in_len {
1377                return Err(SpectrogramError::dimension_mismatch(
1378                    self.inner.in_len,
1379                    input.len(),
1380                ));
1381            }
1382            if output.len() != self.inner.n_fft {
1383                return Err(SpectrogramError::dimension_mismatch(
1384                    self.inner.n_fft,
1385                    output.len(),
1386                ));
1387            }
1388
1389            unsafe {
1390                // Copy input to FFTW buffer
1391                for i in 0..self.inner.in_len {
1392                    let ptr = self.inner.input.as_ptr().add(i).cast::<f64>();
1393                    *ptr.add(0) = input[i].re;
1394                    *ptr.add(1) = input[i].im;
1395                }
1396
1397                // Execute inverse FFT
1398                fftw_sys::fftw_execute_dft_c2r(
1399                    self.inner.plan,
1400                    self.inner.input.as_ptr(),
1401                    self.inner.output.as_ptr(),
1402                );
1403
1404                // Copy output and normalize
1405                let scale = 1.0 / self.inner.n_fft as f64;
1406                for i in 0..self.inner.n_fft {
1407                    output[i] = *self.inner.output.as_ptr().add(i) * scale;
1408                }
1409            }
1410
1411            Ok(())
1412        }
1413    }
1414
1415    impl C2rPlanner for FftwPlanner {
1416        type Plan = FftwInversePlan;
1417
1418        fn plan_c2r(&mut self, n_fft: usize) -> SpectrogramResult<Self::Plan> {
1419            Ok(FftwInversePlan {
1420                inner: self.get_or_create_inverse(n_fft)?,
1421            })
1422        }
1423    }
1424
1425    // ========================================================================
1426    // Complex-to-Complex FFT Implementation
1427    // ========================================================================
1428
1429    /// Inner state for an FFTW complex-to-complex plan pair (forward + inverse).
1430    struct FftwC2cInner {
1431        n_fft: usize,
1432        fwd: fftw_sys::fftw_plan,
1433        inv: fftw_sys::fftw_plan,
1434        buf: Arc<FftwBuffer<fftw_sys::fftw_complex>>,
1435    }
1436
1437    impl Drop for FftwC2cInner {
1438        fn drop(&mut self) {
1439            unsafe {
1440                fftw_sys::fftw_destroy_plan(self.fwd);
1441                fftw_sys::fftw_destroy_plan(self.inv);
1442            }
1443        }
1444    }
1445
1446    /// Complex-to-complex FFT plan (f64) backed by FFTW3.
1447    ///
1448    /// Both directions reuse the same FFTW-aligned buffer; data is copied in/out
1449    /// around each execute call (same pattern as `FftwPlan`).
1450    pub struct FftwC2cPlan {
1451        inner: Arc<FftwC2cInner>,
1452    }
1453
1454    // SAFETY: FFTW plans are safe to use from one thread at a time; FftwC2cPlan is not Clone.
1455    unsafe impl Send for FftwC2cPlan {}
1456
1457    impl FftwPlanner {
1458        /// Build a C2c plan pair for size `n_fft`.
1459        pub(crate) fn build_c2c_plan(n_fft: usize) -> SpectrogramResult<FftwC2cInner> {
1460            let buf = Arc::new(FftwBuffer::<fftw_sys::fftw_complex>::allocate(n_fft)?);
1461            let ptr = buf.as_ptr();
1462
1463            let n_i32: i32 = n_fft
1464                .try_into()
1465                .map_err(|_| SpectrogramError::invalid_input("n_fft too large for FFTW"))?;
1466
1467            let _lock = FFTW_PLANNER_LOCK.lock().map_err(|e| {
1468                SpectrogramError::fft_backend("fftw", format!("FFT planner mutex poisoned: {}", e))
1469            })?;
1470
1471            let fwd = unsafe {
1472                fftw_sys::fftw_plan_dft_1d(
1473                    n_i32,
1474                    ptr,
1475                    ptr,
1476                    fftw_sys::FFTW_FORWARD as i32,
1477                    fftw_sys::FFTW_ESTIMATE,
1478                )
1479            };
1480            if fwd.is_null() {
1481                return Err(SpectrogramError::fft_backend(
1482                    "fftw",
1483                    "failed to create FFTW C2c forward plan",
1484                ));
1485            }
1486
1487            let inv = unsafe {
1488                fftw_sys::fftw_plan_dft_1d(
1489                    n_i32,
1490                    ptr,
1491                    ptr,
1492                    fftw_sys::FFTW_BACKWARD as i32,
1493                    fftw_sys::FFTW_ESTIMATE,
1494                )
1495            };
1496            if inv.is_null() {
1497                unsafe { fftw_sys::fftw_destroy_plan(fwd) };
1498                return Err(SpectrogramError::fft_backend(
1499                    "fftw",
1500                    "failed to create FFTW C2c inverse plan",
1501                ));
1502            }
1503
1504            Ok(FftwC2cInner {
1505                n_fft,
1506                fwd,
1507                inv,
1508                buf,
1509            })
1510        }
1511
1512        /// Get or build a `FftwC2cPlan` (not cached — C2c plans are uncommon enough).
1513        pub fn plan_c2c(&mut self, n_fft: usize) -> SpectrogramResult<FftwC2cPlan> {
1514            Ok(FftwC2cPlan {
1515                inner: Arc::new(Self::build_c2c_plan(n_fft)?),
1516            })
1517        }
1518    }
1519
1520    impl crate::fft_backend::C2cPlan for FftwC2cPlan {
1521        fn n_fft(&self) -> usize {
1522            self.inner.n_fft
1523        }
1524
1525        fn forward(&mut self, buf: &mut [Complex<f64>]) -> SpectrogramResult<()> {
1526            let n = self.inner.n_fft;
1527            if buf.len() != n {
1528                return Err(SpectrogramError::dimension_mismatch(n, buf.len()));
1529            }
1530            unsafe {
1531                let ptr = self.inner.buf.as_ptr();
1532                for i in 0..n {
1533                    let dst = ptr.add(i).cast::<f64>();
1534                    *dst = buf[i].re;
1535                    *dst.add(1) = buf[i].im;
1536                }
1537                fftw_sys::fftw_execute_dft(self.inner.fwd, ptr, ptr);
1538                for i in 0..n {
1539                    let src = ptr.add(i).cast::<f64>();
1540                    buf[i] = Complex::new(*src, *src.add(1));
1541                }
1542            }
1543            Ok(())
1544        }
1545
1546        fn inverse(&mut self, buf: &mut [Complex<f64>]) -> SpectrogramResult<()> {
1547            let n = self.inner.n_fft;
1548            if buf.len() != n {
1549                return Err(SpectrogramError::dimension_mismatch(n, buf.len()));
1550            }
1551            unsafe {
1552                let ptr = self.inner.buf.as_ptr();
1553                for i in 0..n {
1554                    let dst = ptr.add(i).cast::<f64>();
1555                    *dst = buf[i].re;
1556                    *dst.add(1) = buf[i].im;
1557                }
1558                fftw_sys::fftw_execute_dft(self.inner.inv, ptr, ptr);
1559                for i in 0..n {
1560                    let src = ptr.add(i).cast::<f64>();
1561                    buf[i] = Complex::new(*src, *src.add(1));
1562                }
1563            }
1564            Ok(())
1565        }
1566    }
1567
1568    // ========================================================================
1569    // 2D FFT Implementation
1570    // ========================================================================
1571
1572    #[derive(Debug)]
1573    pub(crate) struct PlanInner2d {
1574        nrows: usize,
1575        ncols: usize,
1576        out_shape: (usize, usize),
1577        plan: fftw_sys::fftw_plan,
1578        input: Arc<FftwBuffer<f64>>,
1579        output: Arc<FftwBuffer<fftw_sys::fftw_complex>>,
1580    }
1581
1582    impl Drop for PlanInner2d {
1583        fn drop(&mut self) {
1584            unsafe {
1585                fftw_sys::fftw_destroy_plan(self.plan);
1586            }
1587        }
1588    }
1589
1590    #[derive(Debug)]
1591    pub(crate) struct InversePlanInner2d {
1592        nrows: usize,
1593        ncols: usize,
1594        in_shape: (usize, usize),
1595        plan: fftw_sys::fftw_plan,
1596        input: Arc<FftwBuffer<fftw_sys::fftw_complex>>,
1597        output: Arc<FftwBuffer<f64>>,
1598    }
1599
1600    impl Drop for InversePlanInner2d {
1601        fn drop(&mut self) {
1602            unsafe {
1603                fftw_sys::fftw_destroy_plan(self.plan);
1604            }
1605        }
1606    }
1607
1608    impl FftwPlanner {
1609        pub(crate) fn build_plan_2d(nrows: usize, ncols: usize) -> SpectrogramResult<PlanInner2d> {
1610            let out_shape = crate::fft_backend::r2c_output_size_2d(nrows, ncols);
1611            let input_len = nrows * ncols;
1612            let output_len = out_shape.0 * out_shape.1;
1613
1614            let input = Arc::new(FftwBuffer::<f64>::allocate(input_len)?);
1615            let output = Arc::new(FftwBuffer::<fftw_sys::fftw_complex>::allocate(output_len)?);
1616
1617            let n0: i32 = nrows
1618                .try_into()
1619                .map_err(|_| SpectrogramError::invalid_input("nrows too large for FFTW"))?;
1620            let n1: i32 = ncols
1621                .try_into()
1622                .map_err(|_| SpectrogramError::invalid_input("ncols too large for FFTW"))?;
1623
1624            // FFTW plan creation is not thread-safe - must be serialized
1625            let _lock = FFTW_PLANNER_LOCK.lock().map_err(|e| {
1626                SpectrogramError::fft_backend("fftw", format!("FFT planner mutex poisoned: {}", e))
1627            })?;
1628
1629            let plan = unsafe {
1630                fftw_sys::fftw_plan_dft_r2c_2d(
1631                    n0,
1632                    n1,
1633                    input.as_ptr(),
1634                    output.as_ptr(),
1635                    fftw_sys::FFTW_ESTIMATE,
1636                )
1637            };
1638
1639            if plan.is_null() {
1640                return Err(SpectrogramError::fft_backend(
1641                    "fftw",
1642                    "failed to create FFTW 2D plan",
1643                ));
1644            }
1645
1646            Ok(PlanInner2d {
1647                nrows,
1648                ncols,
1649                out_shape,
1650                plan,
1651                input,
1652                output,
1653            })
1654        }
1655
1656        pub(crate) fn get_or_create_2d(
1657            &mut self,
1658            nrows: usize,
1659            ncols: usize,
1660        ) -> SpectrogramResult<Arc<PlanInner2d>> {
1661            // For 2D plans, we need a different caching strategy
1662            // For now, we'll create a new plan each time
1663            // TODO: Add HashMap<(usize, usize), Arc<PlanInner2d>> cache
1664            let p = Arc::new(Self::build_plan_2d(nrows, ncols)?);
1665            Ok(p)
1666        }
1667
1668        pub(crate) fn build_inverse_plan_2d(
1669            nrows: usize,
1670            ncols: usize,
1671        ) -> SpectrogramResult<InversePlanInner2d> {
1672            let in_shape = crate::fft_backend::r2c_output_size_2d(nrows, ncols);
1673            let input_len = in_shape.0 * in_shape.1;
1674            let output_len = nrows * ncols;
1675
1676            let input = Arc::new(FftwBuffer::<fftw_sys::fftw_complex>::allocate(input_len)?);
1677            let output = Arc::new(FftwBuffer::<f64>::allocate(output_len)?);
1678
1679            let n0: i32 = nrows
1680                .try_into()
1681                .map_err(|_| SpectrogramError::invalid_input("nrows too large for FFTW"))?;
1682            let n1: i32 = ncols
1683                .try_into()
1684                .map_err(|_| SpectrogramError::invalid_input("ncols too large for FFTW"))?;
1685
1686            // FFTW plan creation is not thread-safe - must be serialized
1687            let _lock = FFTW_PLANNER_LOCK.lock().map_err(|e| {
1688                SpectrogramError::fft_backend("fftw", format!("FFT planner mutex poisoned: {}", e))
1689            })?;
1690
1691            let plan = unsafe {
1692                fftw_sys::fftw_plan_dft_c2r_2d(
1693                    n0,
1694                    n1,
1695                    input.as_ptr(),
1696                    output.as_ptr(),
1697                    fftw_sys::FFTW_ESTIMATE,
1698                )
1699            };
1700
1701            if plan.is_null() {
1702                return Err(SpectrogramError::fft_backend(
1703                    "fftw",
1704                    "failed to create FFTW 2D inverse plan",
1705                ));
1706            }
1707
1708            Ok(InversePlanInner2d {
1709                nrows,
1710                ncols,
1711                in_shape,
1712                plan,
1713                input,
1714                output,
1715            })
1716        }
1717
1718        pub(crate) fn get_or_create_inverse_2d(
1719            &mut self,
1720            nrows: usize,
1721            ncols: usize,
1722        ) -> SpectrogramResult<Arc<InversePlanInner2d>> {
1723            // For 2D plans, we need a different caching strategy
1724            // For now, we'll create a new plan each time
1725            // TODO: Add HashMap<(usize, usize), Arc<InversePlanInner2d>> cache
1726            let p = Arc::new(Self::build_inverse_plan_2d(nrows, ncols)?);
1727            Ok(p)
1728        }
1729    }
1730
1731    #[derive(Debug, Clone)]
1732    pub struct FftwPlan2d {
1733        inner: Arc<PlanInner2d>,
1734    }
1735
1736    impl crate::fft_backend::R2cPlan2d for FftwPlan2d {
1737        fn process(&mut self, input: &[f64], output: &mut [Complex<f64>]) -> SpectrogramResult<()> {
1738            crate::fft_backend::validate_fft_io_2d(
1739                self.inner.nrows,
1740                self.inner.ncols,
1741                input,
1742                output,
1743            )?;
1744
1745            unsafe {
1746                // Copy input to FFTW buffer
1747                core::ptr::copy_nonoverlapping(
1748                    input.as_ptr(),
1749                    self.inner.input.as_ptr(),
1750                    input.len(),
1751                );
1752
1753                // Execute 2D FFT
1754                fftw_sys::fftw_execute_dft_r2c(
1755                    self.inner.plan,
1756                    self.inner.input.as_ptr(),
1757                    self.inner.output.as_ptr(),
1758                );
1759
1760                // Copy output (convert from fftw_complex to Complex<f64>)
1761                let output_len = self.inner.out_shape.0 * self.inner.out_shape.1;
1762                for i in 0..output_len {
1763                    let c = self.inner.output.as_ptr().add(i) as *const f64;
1764                    let re = *c.add(0);
1765                    let im = *c.add(1);
1766                    output[i] = Complex::new(re, im);
1767                }
1768            }
1769
1770            Ok(())
1771        }
1772    }
1773
1774    impl crate::fft_backend::R2cPlanner2d for FftwPlanner {
1775        type Plan = FftwPlan2d;
1776
1777        fn plan_r2c_2d(&mut self, nrows: usize, ncols: usize) -> SpectrogramResult<Self::Plan> {
1778            Ok(FftwPlan2d {
1779                inner: self.get_or_create_2d(nrows, ncols)?,
1780            })
1781        }
1782    }
1783
1784    #[derive(Debug, Clone)]
1785    pub struct FftwInversePlan2d {
1786        inner: Arc<InversePlanInner2d>,
1787    }
1788
1789    impl crate::fft_backend::C2rPlan2d for FftwInversePlan2d {
1790        fn process(&mut self, input: &[Complex<f64>], output: &mut [f64]) -> SpectrogramResult<()> {
1791            let expected_in_len = self.inner.in_shape.0 * self.inner.in_shape.1;
1792            if input.len() != expected_in_len {
1793                return Err(SpectrogramError::dimension_mismatch(
1794                    expected_in_len,
1795                    input.len(),
1796                ));
1797            }
1798
1799            let expected_out_len = self.inner.nrows * self.inner.ncols;
1800            if output.len() != expected_out_len {
1801                return Err(SpectrogramError::dimension_mismatch(
1802                    expected_out_len,
1803                    output.len(),
1804                ));
1805            }
1806
1807            unsafe {
1808                // Copy input to FFTW buffer
1809                for i in 0..input.len() {
1810                    let ptr = self.inner.input.as_ptr().add(i).cast::<f64>();
1811                    *ptr.add(0) = input[i].re;
1812                    *ptr.add(1) = input[i].im;
1813                }
1814
1815                // Execute inverse 2D FFT
1816                fftw_sys::fftw_execute_dft_c2r(
1817                    self.inner.plan,
1818                    self.inner.input.as_ptr(),
1819                    self.inner.output.as_ptr(),
1820                );
1821
1822                // Copy output and normalize
1823                let scale = 1.0 / (self.inner.nrows * self.inner.ncols) as f64;
1824                for i in 0..expected_out_len {
1825                    output[i] = *self.inner.output.as_ptr().add(i) * scale;
1826                }
1827            }
1828
1829            Ok(())
1830        }
1831    }
1832
1833    impl crate::fft_backend::C2rPlanner2d for FftwPlanner {
1834        type Plan = FftwInversePlan2d;
1835
1836        fn plan_c2r_2d(&mut self, nrows: usize, ncols: usize) -> SpectrogramResult<Self::Plan> {
1837            Ok(FftwInversePlan2d {
1838                inner: self.get_or_create_inverse_2d(nrows, ncols)?,
1839            })
1840        }
1841    }
1842}
1843
1844#[cfg(all(test, feature = "realfft"))]
1845mod tests_c2c {
1846    use super::realfft_backend::{RealFftC2cPlan, RealFftC2cPlanF32};
1847    use crate::fft_backend::{C2cPlan, C2cPlanF32};
1848    use num_complex::Complex;
1849
1850    /// DFT of a pure DC signal (all ones): X[0] = N, X[k≠0] = 0.
1851    #[test]
1852    fn c2c_f64_dc_signal() {
1853        let n = 8;
1854        let mut plan = RealFftC2cPlan::new(n);
1855        let mut buf: Vec<Complex<f64>> = vec![Complex::new(1.0, 0.0); n];
1856        plan.forward(&mut buf).unwrap();
1857        assert!((buf[0].re - n as f64).abs() < 1e-10, "DC bin should be N");
1858        assert!(buf[0].im.abs() < 1e-10);
1859        for k in 1..n {
1860            assert!(buf[k].norm() < 1e-10, "bin {k} should be zero");
1861        }
1862    }
1863
1864    /// Round-trip: forward then inverse (divided by N) recovers the original signal.
1865    #[test]
1866    fn c2c_f64_round_trip() {
1867        let n = 16;
1868        let original: Vec<Complex<f64>> = (0..n)
1869            .map(|i| Complex::new(i as f64, -(i as f64)))
1870            .collect();
1871        let mut buf = original.clone();
1872        let mut plan = RealFftC2cPlan::new(n);
1873        plan.forward(&mut buf).unwrap();
1874        plan.inverse(&mut buf).unwrap();
1875        let scale = 1.0 / n as f64;
1876        for (a, b) in buf.iter().zip(original.iter()) {
1877            assert!((a.re * scale - b.re).abs() < 1e-10);
1878            assert!((a.im * scale - b.im).abs() < 1e-10);
1879        }
1880    }
1881
1882    /// f32 round-trip.
1883    #[test]
1884    fn c2c_f32_round_trip() {
1885        let n = 16;
1886        let original: Vec<Complex<f32>> = (0..n)
1887            .map(|i| Complex::new(i as f32, -(i as f32)))
1888            .collect();
1889        let mut buf = original.clone();
1890        let mut plan = RealFftC2cPlanF32::new(n);
1891        plan.forward(&mut buf).unwrap();
1892        plan.inverse(&mut buf).unwrap();
1893        let scale = 1.0f32 / n as f32;
1894        for (a, b) in buf.iter().zip(original.iter()) {
1895            assert!((a.re * scale - b.re).abs() < 1e-5);
1896            assert!((a.im * scale - b.im).abs() < 1e-5);
1897        }
1898    }
1899
1900    /// Dimension mismatch returns an error rather than panicking.
1901    #[test]
1902    fn c2c_f64_dimension_mismatch() {
1903        let mut plan = RealFftC2cPlan::new(8);
1904        let mut buf = vec![Complex::new(0.0, 0.0); 7]; // wrong size
1905        assert!(plan.forward(&mut buf).is_err());
1906        assert!(plan.inverse(&mut buf).is_err());
1907    }
1908}
spectrograms/fft_backend.rs

spectrograms/
fft_backend.rs
