Skip to main content

spectrograms/
spectrogram.rs

1use std::marker::PhantomData;
2use std::num::NonZeroUsize;
3use std::ops::{Deref, DerefMut};
4
5use ndarray::{Array1, Array2};
6use non_empty_slice::{NonEmptySlice, NonEmptyVec, non_empty_vec};
7use num_complex::Complex;
8
9#[cfg(feature = "python")]
10use pyo3::prelude::*;
11
12use crate::cqt::CqtKernel;
13use crate::erb::ErbFilterbank;
14use crate::{
15    CqtParams, ErbParams, R2cPlan, Sample, SpectrogramError, SpectrogramResult, WindowType,
16    min_max_single_pass, nzu,
17};
18const EPS: f64 = 1e-12;
19
20//
21// ========================
22// Sparse Matrix for efficient filterbank multiplication
23// ========================
24//
25
26/// Row-wise sparse matrix optimized for matrix-vector multiplication.
27///
28/// This structure stores sparse data as a vector of vectors, where each row maintains
29/// its own list of non-zero values and corresponding column indices. This is more
30/// flexible than traditional CSR format and allows efficient row-by-row construction.
31///
32/// This structure is designed for matrices with very few non-zero values per row,
33/// such as mel filterbanks (triangular filters) and logarithmic frequency mappings
34/// (linear interpolation between 1-2 adjacent bins).
35///
36/// For typical spectrograms:
37/// - `LogHz` interpolation: Only 1-2 non-zeros per row (~99% sparse)
38/// - Mel filterbank: ~10-50 non-zeros per row depending on FFT size (~90-98% sparse)
39///
40/// By storing only non-zero values, we avoid wasting CPU cycles multiplying by zero,
41/// which can provide 10-100x speedup compared to dense matrix multiplication.
42#[derive(Debug, Clone)]
43struct SparseMatrix {
44    /// Number of rows
45    nrows: usize,
46    /// Number of columns
47    ncols: usize,
48    /// Non-zero values for each row (row-major order)
49    values: Vec<Vec<f64>>,
50    /// Column indices for each non-zero value
51    indices: Vec<Vec<usize>>,
52}
53
54impl SparseMatrix {
55    /// Create a new sparse matrix with the given dimensions.
56    fn new(nrows: usize, ncols: usize) -> Self {
57        Self {
58            nrows,
59            ncols,
60            values: vec![Vec::new(); nrows],
61            indices: vec![Vec::new(); nrows],
62        }
63    }
64
65    /// Set a value in the matrix. Only stores if value is non-zero.
66    ///
67    /// # Panics (debug mode only)
68    /// Panics in debug builds if row or col are out of bounds.
69    fn set(&mut self, row: usize, col: usize, value: f64) {
70        debug_assert!(
71            row < self.nrows && col < self.ncols,
72            "SparseMatrix index out of bounds: ({}, {}) for {}x{} matrix",
73            row,
74            col,
75            self.nrows,
76            self.ncols
77        );
78        if row >= self.nrows || col >= self.ncols {
79            return;
80        }
81
82        // Only store non-zero values (with small epsilon for numerical stability)
83        if value.abs() > 1e-10 {
84            self.values[row].push(value);
85            self.indices[row].push(col);
86        }
87    }
88
89    /// Get the number of rows.
90    const fn nrows(&self) -> usize {
91        self.nrows
92    }
93
94    /// Get the number of columns.
95    const fn ncols(&self) -> usize {
96        self.ncols
97    }
98
99    /// Perform sparse matrix-vector multiplication: out = self * input
100    /// This is much faster than dense multiplication when the matrix is sparse.
101    #[inline]
102    fn multiply_vec<T: Sample>(&self, input: &[T], out: &mut [T]) {
103        debug_assert_eq!(input.len(), self.ncols);
104        debug_assert_eq!(out.len(), self.nrows);
105
106        // Filterbank coefficients are constructed in f64 (unchanged math) and
107        // converted to `T` at apply time, where they multiply the `T` data.
108        for (row_idx, (row_values, row_indices)) in
109            self.values.iter().zip(&self.indices).enumerate()
110        {
111            let mut acc = T::zero();
112            for (&value, &col_idx) in row_values.iter().zip(row_indices) {
113                acc += T::from_f64(value) * input[col_idx];
114            }
115            out[row_idx] = acc;
116        }
117    }
118}
119
120// Linear frequency
121pub type LinearPowerSpectrogram = Spectrogram<LinearHz, Power>;
122pub type LinearMagnitudeSpectrogram = Spectrogram<LinearHz, Magnitude>;
123pub type LinearDbSpectrogram = Spectrogram<LinearHz, Decibels>;
124pub type LinearSpectrogram<AmpScale> = Spectrogram<LinearHz, AmpScale>;
125
126// Log-frequency (e.g. CQT-style)
127pub type LogHzPowerSpectrogram = Spectrogram<LogHz, Power>;
128pub type LogHzMagnitudeSpectrogram = Spectrogram<LogHz, Magnitude>;
129pub type LogHzDbSpectrogram = Spectrogram<LogHz, Decibels>;
130pub type LogHzSpectrogram<AmpScale> = Spectrogram<LogHz, AmpScale>;
131
132// ERB / gammatone
133pub type ErbPowerSpectrogram = Spectrogram<Erb, Power>;
134pub type ErbMagnitudeSpectrogram = Spectrogram<Erb, Magnitude>;
135pub type ErbDbSpectrogram = Spectrogram<Erb, Decibels>;
136pub type GammatonePowerSpectrogram = ErbPowerSpectrogram;
137pub type GammatoneMagnitudeSpectrogram = ErbMagnitudeSpectrogram;
138pub type GammatoneDbSpectrogram = ErbDbSpectrogram;
139pub type ErbSpectrogram<AmpScale> = Spectrogram<Erb, AmpScale>;
140pub type GammatoneSpectrogram<AmpScale> = ErbSpectrogram<AmpScale>;
141
142// Mel
143pub type MelMagnitudeSpectrogram = Spectrogram<Mel, Magnitude>;
144pub type MelPowerSpectrogram = Spectrogram<Mel, Power>;
145pub type MelDbSpectrogram = Spectrogram<Mel, Decibels>;
146pub type LogMelSpectrogram = MelDbSpectrogram;
147pub type MelSpectrogram<AmpScale> = Spectrogram<Mel, AmpScale>;
148
149// CQT
150pub type CqtPowerSpectrogram = Spectrogram<Cqt, Power>;
151pub type CqtMagnitudeSpectrogram = Spectrogram<Cqt, Magnitude>;
152pub type CqtDbSpectrogram = Spectrogram<Cqt, Decibels>;
153pub type CqtSpectrogram<AmpScale> = Spectrogram<Cqt, AmpScale>;
154
155use crate::fft_backend::r2c_output_size;
156
157/// A spectrogram plan is the compiled, reusable execution object.
158///
159/// It owns:
160/// - FFT plan (reusable)
161/// - window samples
162/// - mapping (identity / mel filterbank / etc.)
163/// - amplitude scaling config
164/// - workspace buffers to avoid allocations in hot loops
165///
166/// It computes one specific spectrogram type: `Spectrogram<FreqScale, AmpScale>`.
167///
168/// # Type Parameters
169///
170/// - `FreqScale`: Frequency scale type (e.g. `LinearHz`, `LogHz`, `Mel`, etc.)
171/// - `AmpScale`: Amplitude scaling type (e.g. `Power`, `Magnitude`, `Decibels`, etc.)
172pub struct SpectrogramPlan<FreqScale, AmpScale, T = f64>
173where
174    AmpScale: AmpScaleSpec + 'static,
175    FreqScale: Copy + Clone + 'static,
176{
177    params: SpectrogramParams,
178
179    stft: StftPlan<T>,
180    mapping: FrequencyMapping<FreqScale>,
181    scaling: AmplitudeScaling<AmpScale>,
182
183    freq_axis: FrequencyAxis<FreqScale>,
184    workspace: Workspace<T>,
185
186    _amp: PhantomData<AmpScale>,
187}
188
189impl<FreqScale, AmpScale, T> SpectrogramPlan<FreqScale, AmpScale, T>
190where
191    AmpScale: AmpScaleSpec + 'static,
192    FreqScale: Copy + Clone + 'static,
193    T: Sample,
194{
195    /// Get the spectrogram parameters used to create this plan.
196    ///
197    /// # Returns
198    ///
199    /// A reference to the `SpectrogramParams` used in this plan.
200    #[inline]
201    #[must_use]
202    pub const fn params(&self) -> &SpectrogramParams {
203        &self.params
204    }
205
206    /// Get the frequency axis for this spectrogram plan.
207    ///
208    /// # Returns
209    ///
210    /// A reference to the `FrequencyAxis<FreqScale>` used in this plan.
211    #[inline]
212    #[must_use]
213    pub const fn freq_axis(&self) -> &FrequencyAxis<FreqScale> {
214        &self.freq_axis
215    }
216
217    /// Compute a spectrogram for a mono signal.
218    ///
219    /// This function performs:
220    /// - framing + windowing
221    /// - FFT per frame
222    /// - magnitude/power
223    /// - frequency mapping (identity/mel/etc.)
224    /// - amplitude scaling (linear or dB)
225    ///
226    /// It allocates the output `Array2` once, but does not allocate per-frame.
227    ///
228    /// # Arguments
229    ///
230    /// * `samples` - Audio samples
231    ///
232    /// # Returns
233    ///
234    /// A `Spectrogram<FreqScale, AmpScale>` containing the computed spectrogram.
235    ///
236    /// # Errors
237    ///
238    /// Returns an error if STFT computation or mapping fails.
239    #[inline]
240    pub fn compute(
241        &mut self,
242        samples: &NonEmptySlice<T>,
243    ) -> SpectrogramResult<Spectrogram<FreqScale, AmpScale, T>> {
244        let n_frames = self.stft.frame_count(samples.len())?;
245        let n_bins = self.mapping.output_bins();
246
247        // Create output matrix: (n_bins, n_frames)
248        let mut data = Array2::<T>::zeros((n_bins.get(), n_frames.get()));
249
250        // Ensure workspace is correctly sized
251        self.workspace
252            .ensure_sizes(self.stft.n_fft, self.stft.out_len, n_bins);
253
254        // Main loop: fill each frame (column)
255        for frame_idx in 0..n_frames.get() {
256            // CQT needs unwindowed frames because its kernels already contain windowing.
257            // Other mappings use the FFT spectrum, so they need the windowed frame.
258            if self.mapping.kind.needs_unwindowed_frame() {
259                // Fill frame without windowing (for CQT)
260                self.stft
261                    .fill_frame_unwindowed(samples, frame_idx, &mut self.workspace)?;
262            } else {
263                // Compute windowed frame spectrum (for FFT-based mappings)
264                self.stft
265                    .compute_frame_spectrum(samples, frame_idx, &mut self.workspace)?;
266            }
267
268            // mapping: spectrum(out_len) -> mapped(n_bins)
269            // For CQT, this uses workspace.frame; for others, workspace.spectrum
270            // For ERB, we need the complex FFT output (fft_out)
271            // We need to borrow workspace fields separately to avoid borrow conflicts
272            let Workspace {
273                spectrum,
274                mapped,
275                frame,
276                ..
277            } = &mut self.workspace;
278
279            self.mapping.apply(spectrum, frame, mapped)?;
280
281            // amplitude scaling in-place on mapped vector
282            self.scaling.apply_in_place(mapped)?;
283
284            // write column into output
285            for (row, &val) in mapped.iter().enumerate() {
286                data[[row, frame_idx]] = val;
287            }
288        }
289
290        let times = build_time_axis_seconds(&self.params, n_frames);
291        let axes = Axes::new(self.freq_axis.clone(), times);
292
293        Ok(Spectrogram::new(data, axes, self.params.clone()))
294    }
295
296    /// Compute a single frame of the spectrogram.
297    ///
298    /// This is useful for streaming/online processing where you want to
299    /// process audio frame-by-frame without computing the entire spectrogram.
300    ///
301    /// # Arguments
302    ///
303    /// * `samples` - Audio samples (must contain at least enough samples for the requested frame)
304    /// * `frame_idx` - Frame index to compute
305    ///
306    /// # Returns
307    ///
308    /// A vector of frequency bin values for the requested frame.
309    ///
310    /// # Errors
311    ///
312    /// Returns an error if the frame index is out of bounds or if STFT computation fails.
313    ///
314    /// # Examples
315    ///
316    /// ```
317    /// use spectrograms::*;
318    /// use non_empty_slice::non_empty_vec;
319    ///
320    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
321    /// let samples = non_empty_vec![0.0; nzu!(16000)];
322    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
323    /// let params = SpectrogramParams::new(stft, 16000.0)?;
324    ///
325    /// let planner = SpectrogramPlanner::new();
326    /// let mut plan = planner.linear_plan::<Power, _>(&params, None)?;
327    ///
328    /// // Compute just the first frame
329    /// let frame = plan.compute_frame(&samples, 0)?;
330    /// assert_eq!(frame.len(), nzu!(257)); // n_fft/2 + 1
331    /// # Ok(())
332    /// # }
333    /// ```
334    #[inline]
335    pub fn compute_frame(
336        &mut self,
337        samples: &NonEmptySlice<T>,
338        frame_idx: usize,
339    ) -> SpectrogramResult<NonEmptyVec<T>> {
340        let n_bins = self.mapping.output_bins();
341
342        // Ensure workspace is correctly sized
343        self.workspace
344            .ensure_sizes(self.stft.n_fft, self.stft.out_len, n_bins);
345
346        // CQT needs unwindowed frames because its kernels already contain windowing.
347        // Other mappings use the FFT spectrum, so they need the windowed frame.
348        if self.mapping.kind.needs_unwindowed_frame() {
349            // Fill frame without windowing (for CQT)
350            self.stft
351                .fill_frame_unwindowed(samples, frame_idx, &mut self.workspace)?;
352        } else {
353            // Compute windowed frame spectrum (for FFT-based mappings)
354            self.stft
355                .compute_frame_spectrum(samples, frame_idx, &mut self.workspace)?;
356        }
357
358        // Apply mapping (using split borrows to avoid borrow conflicts)
359        let Workspace {
360            spectrum,
361            mapped,
362            frame,
363            ..
364        } = &mut self.workspace;
365
366        self.mapping.apply(spectrum, frame, mapped)?;
367
368        // Apply amplitude scaling
369        self.scaling.apply_in_place(mapped)?;
370
371        Ok(mapped.clone())
372    }
373
374    /// Compute spectrogram into a pre-allocated buffer.
375    ///
376    /// This avoids allocating the output matrix, which is useful when
377    /// you want to reuse buffers or have strict memory requirements.
378    ///
379    /// # Arguments
380    ///
381    /// * `samples` - Audio samples
382    /// * `output` - Pre-allocated output matrix (must be correct size: `n_bins` x `n_frames`)
383    ///
384    /// # Returns
385    ///
386    /// An empty result on success.
387    ///
388    /// # Errors
389    ///
390    /// Returns an error if the output buffer dimensions don't match the expected size.
391    ///
392    /// # Examples
393    ///
394    /// ```
395    /// use spectrograms::*;
396    /// use ndarray::Array2;
397    /// use non_empty_slice::non_empty_vec;
398    ///
399    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
400    /// let samples = non_empty_vec![0.0; nzu!(16000)];
401    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
402    /// let params = SpectrogramParams::new(stft, 16000.0)?;
403    ///
404    /// let planner = SpectrogramPlanner::new();
405    /// let mut plan = planner.linear_plan::<Power, _>(&params, None)?;
406    ///
407    /// // Pre-allocate output buffer
408    /// let mut output = Array2::<f64>::zeros((257, 63));
409    /// plan.compute_into(&samples, &mut output)?;
410    /// # Ok(())
411    /// # }
412    /// ```
413    #[inline]
414    pub fn compute_into(
415        &mut self,
416        samples: &NonEmptySlice<T>,
417        output: &mut Array2<T>,
418    ) -> SpectrogramResult<()> {
419        let n_frames = self.stft.frame_count(samples.len())?;
420        let n_bins = self.mapping.output_bins();
421
422        // Validate output dimensions
423        if output.nrows() != n_bins.get() {
424            return Err(SpectrogramError::dimension_mismatch(
425                n_bins.get(),
426                output.nrows(),
427            ));
428        }
429        if output.ncols() != n_frames.get() {
430            return Err(SpectrogramError::dimension_mismatch(
431                n_frames.get(),
432                output.ncols(),
433            ));
434        }
435
436        // Ensure workspace is correctly sized
437        self.workspace
438            .ensure_sizes(self.stft.n_fft, self.stft.out_len, n_bins);
439
440        // Main loop: fill each frame (column)
441        for frame_idx in 0..n_frames.get() {
442            // CQT needs unwindowed frames because its kernels already contain windowing.
443            // Other mappings use the FFT spectrum, so they need the windowed frame.
444            if self.mapping.kind.needs_unwindowed_frame() {
445                // Fill frame without windowing (for CQT)
446                self.stft
447                    .fill_frame_unwindowed(samples, frame_idx, &mut self.workspace)?;
448            } else {
449                // Compute windowed frame spectrum (for FFT-based mappings)
450                self.stft
451                    .compute_frame_spectrum(samples, frame_idx, &mut self.workspace)?;
452            }
453
454            // mapping: spectrum(out_len) -> mapped(n_bins)
455            // For CQT, this uses workspace.frame; for others, workspace.spectrum
456            // For ERB, we need the complex FFT output (fft_out)
457            // We need to borrow workspace fields separately to avoid borrow conflicts
458            let Workspace {
459                spectrum,
460                mapped,
461                frame,
462                ..
463            } = &mut self.workspace;
464
465            self.mapping.apply(spectrum, frame, mapped)?;
466
467            // amplitude scaling in-place on mapped vector
468            self.scaling.apply_in_place(mapped)?;
469
470            // write column into output
471            for (row, &val) in mapped.iter().enumerate() {
472                output[[row, frame_idx]] = val;
473            }
474        }
475
476        Ok(())
477    }
478
479    /// Get the expected output dimensions for a given signal length.
480    ///
481    /// # Arguments
482    ///
483    /// * `signal_length` - Length of the input signal in samples.
484    ///
485    /// # Returns
486    ///
487    /// A tuple `(n_bins, n_frames)` representing the output spectrogram shape.
488    ///
489    /// # Errors
490    ///
491    /// Returns an error if the signal length is too short to produce any frames.
492    ///
493    /// # Examples
494    ///
495    /// ```
496    /// use spectrograms::*;
497    ///
498    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
499    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
500    /// let params = SpectrogramParams::new(stft, 16000.0)?;
501    ///
502    /// let planner = SpectrogramPlanner::new();
503    /// let plan = planner.linear_plan::<Power, f64>(&params, None)?;
504    ///
505    /// let (n_bins, n_frames) = plan.output_shape(nzu!(16000))?;
506    /// assert_eq!(n_bins, nzu!(257));
507    /// assert_eq!(n_frames, nzu!(63));
508    /// # Ok(())
509    /// # }
510    /// ```
511    #[inline]
512    pub fn output_shape(
513        &self,
514        signal_length: NonZeroUsize,
515    ) -> SpectrogramResult<(NonZeroUsize, NonZeroUsize)> {
516        let n_frames = self.stft.frame_count(signal_length)?;
517        let n_bins = self.mapping.output_bins();
518        Ok((n_bins, n_frames))
519    }
520}
521
522/// STFT (Short-Time Fourier Transform) result containing complex frequency bins.
523///
524/// This is the raw STFT output before any frequency mapping or amplitude scaling.
525///
526/// # Fields
527///
528/// - `data`: Complex STFT matrix with shape (`frequency_bins`, `time_frames`)
529/// - `frequencies`: Frequency axis in Hz
530/// - `sample_rate`: Sample rate in Hz
531/// - `params`: STFT computation parameters
532#[derive(Debug, Clone)]
533#[non_exhaustive]
534pub struct StftResult<T = f64> {
535    /// Complex STFT matrix with shape (`frequency_bins`, `time_frames`)
536    pub data: Array2<Complex<T>>,
537    /// Frequency axis in Hz
538    pub frequencies: NonEmptyVec<f64>,
539    /// Sample rate in Hz
540    pub sample_rate: f64,
541    pub params: StftParams,
542}
543
544impl<T: Sample> StftResult<T> {
545    /// Get the number of frequency bins.
546    ///
547    /// # Returns
548    ///
549    /// Number of frequency bins in the STFT result.
550    #[inline]
551    #[must_use]
552    pub fn n_bins(&self) -> NonZeroUsize {
553        // safety: nrows() > 0 for NonEmptyVec
554        unsafe { NonZeroUsize::new_unchecked(self.data.nrows()) }
555    }
556
557    /// Get the number of time frames.
558    ///
559    /// # Returns
560    ///
561    /// Number of time frames in the STFT result.
562    #[inline]
563    #[must_use]
564    pub fn n_frames(&self) -> NonZeroUsize {
565        // safety: ncols() > 0 for NonEmptyVec
566        unsafe { NonZeroUsize::new_unchecked(self.data.ncols()) }
567    }
568
569    /// Get the frequency resolution in Hz
570    ///
571    /// # Returns
572    ///
573    /// Frequency bin width in Hz.
574    #[inline]
575    #[must_use]
576    pub fn frequency_resolution(&self) -> f64 {
577        self.sample_rate / self.params.n_fft().get() as f64
578    }
579
580    /// Get the time resolution in seconds.
581    ///
582    /// # Returns
583    ///
584    /// Time between successive frames in seconds.
585    #[inline]
586    #[must_use]
587    pub fn time_resolution(&self) -> f64 {
588        self.params.hop_size().get() as f64 / self.sample_rate
589    }
590
591    /// Normalizes self.data to remove the complex aspect of it.
592    ///
593    /// # Returns
594    ///
595    /// An Array2\<f64\> containing the norms of each complex number in self.data.
596    #[inline]
597    pub fn norm(&self) -> Array2<T> {
598        self.as_ref().mapv(Complex::norm)
599    }
600}
601
602impl<T: Sample> AsRef<Array2<Complex<T>>> for StftResult<T> {
603    #[inline]
604    fn as_ref(&self) -> &Array2<Complex<T>> {
605        &self.data
606    }
607}
608
609impl<T: Sample> AsMut<Array2<Complex<T>>> for StftResult<T> {
610    #[inline]
611    fn as_mut(&mut self) -> &mut Array2<Complex<T>> {
612        &mut self.data
613    }
614}
615
616impl<T: Sample> Deref for StftResult<T> {
617    type Target = Array2<Complex<T>>;
618
619    #[inline]
620    fn deref(&self) -> &Self::Target {
621        &self.data
622    }
623}
624
625impl<T: Sample> DerefMut for StftResult<T> {
626    #[inline]
627    fn deref_mut(&mut self) -> &mut Self::Target {
628        &mut self.data
629    }
630}
631
632/// A planner is an object that can build spectrogram plans.
633///
634/// In your design, this is where:
635/// - FFT plans are created
636/// - mapping matrices are compiled
637/// - axes are computed
638///
639/// This allows you to keep plan building separate from the output types.
640#[derive(Debug, Default)]
641#[non_exhaustive]
642pub struct SpectrogramPlanner;
643
644impl SpectrogramPlanner {
645    /// Create a new spectrogram planner.
646    ///
647    /// # Returns
648    ///
649    /// A new `SpectrogramPlanner` instance.
650    #[inline]
651    #[must_use]
652    pub const fn new() -> Self {
653        Self
654    }
655
656    /// Compute the Short-Time Fourier Transform (STFT) of a signal.
657    ///
658    /// This returns the raw complex STFT matrix before any frequency mapping
659    /// or amplitude scaling. Useful for applications that need the full complex
660    /// spectrum or custom processing.
661    ///
662    /// # Arguments
663    ///
664    /// * `samples` - Audio samples (any type that can be converted to a slice)
665    /// * `params` - STFT computation parameters
666    ///
667    /// # Returns
668    ///
669    /// An `StftResult` containing the complex STFT matrix and metadata.
670    ///
671    /// # Errors
672    ///
673    /// Returns an error if STFT computation fails.
674    ///
675    /// # Examples
676    ///
677    /// ```
678    /// use spectrograms::*;
679    /// use non_empty_slice::non_empty_vec;
680    ///
681    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
682    /// let samples = non_empty_vec![0.0; nzu!(16000)];
683    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
684    /// let params = SpectrogramParams::new(stft, 16000.0)?;
685    ///
686    /// let planner = SpectrogramPlanner::new();
687    /// let stft_result = planner.compute_stft(&samples, &params)?;
688    ///
689    /// println!("STFT: {} bins x {} frames", stft_result.n_bins(), stft_result.n_frames());
690    /// # Ok(())
691    /// # }
692    /// ```
693    ///
694    /// # Performance Note
695    ///
696    /// This method creates a new FFT plan each time. For processing multiple
697    /// signals, create a reusable plan with `StftPlan::new()` instead.
698    ///
699    /// # Examples
700    ///
701    /// ```rust
702    /// use spectrograms::*;
703    /// use non_empty_slice::non_empty_vec;
704    ///
705    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
706    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
707    /// let params = SpectrogramParams::new(stft, 16000.0)?;
708    ///
709    /// // One-shot (convenient)
710    /// let planner = SpectrogramPlanner::new();
711    /// let stft_result = planner.compute_stft(&non_empty_vec![0.0; nzu!(16000)], &params)?;
712    ///
713    /// // Reusable plan (efficient for batch)
714    /// let mut plan = StftPlan::new(&params)?;
715    /// for signal in &[non_empty_vec![0.0; nzu!(16000)], non_empty_vec![1.0; nzu!(16000)]] {
716    ///     let stft = plan.compute(&signal, &params)?;
717    /// }
718    /// # Ok(())
719    /// # }
720    /// ```
721    #[inline]
722    pub fn compute_stft<T: Sample>(
723        &self,
724        samples: &NonEmptySlice<T>,
725        params: &SpectrogramParams,
726    ) -> SpectrogramResult<StftResult<T>> {
727        let mut plan = StftPlan::<T>::new(params)?;
728        plan.compute(samples, params)
729    }
730
731    /// Compute the power spectrum of a single audio frame.
732    ///
733    /// This is useful for real-time processing or analyzing individual frames.
734    ///
735    /// # Arguments
736    ///
737    /// * `samples` - Audio frame (length ≤ n_fft, will be zero-padded if shorter)
738    /// * `n_fft` - FFT size
739    /// * `window` - Window type to apply
740    ///
741    /// # Returns
742    ///
743    /// A vector of power values (|X|²) with length `n_fft/2` + 1.
744    ///
745    /// # Automatic Zero-Padding
746    ///
747    /// If the input signal is shorter than `n_fft`, it will be automatically
748    /// zero-padded to the required length.
749    ///
750    /// # Errors
751    ///
752    /// Returns `InvalidInput` error if the input length exceeds `n_fft`.
753    ///
754    /// # Examples
755    ///
756    /// ```
757    /// use spectrograms::*;
758    /// use non_empty_slice::non_empty_vec;
759    ///
760    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
761    /// let frame = non_empty_vec![0.0; nzu!(512)];
762    ///
763    /// let planner = SpectrogramPlanner::new();
764    /// let power = planner.compute_power_spectrum(frame.as_ref(), nzu!(512), WindowType::Hanning)?;
765    ///
766    /// assert_eq!(power.len(), nzu!(257)); // 512/2 + 1
767    /// # Ok(())
768    /// # }
769    /// ```
770    #[inline]
771    pub fn compute_power_spectrum(
772        &self,
773        samples: &NonEmptySlice<f64>,
774        n_fft: NonZeroUsize,
775        window: WindowType,
776    ) -> SpectrogramResult<NonEmptyVec<f64>> {
777        if samples.len() > n_fft {
778            return Err(SpectrogramError::invalid_input(format!(
779                "Input length ({}) exceeds FFT size ({})",
780                samples.len(),
781                n_fft
782            )));
783        }
784
785        let window_samples = make_window::<f64>(window, n_fft);
786        let out_len = r2c_output_size(n_fft.get());
787
788        // Create FFT plan
789        #[cfg(feature = "realfft")]
790        let mut fft = {
791            let mut planner = crate::RealFftPlanner::new();
792            let plan = planner.get_or_create(n_fft.get());
793            crate::RealFftPlan::new(n_fft.get(), plan)
794        };
795
796        #[cfg(feature = "fftw")]
797        let mut fft = {
798            use std::sync::Arc;
799            let plan = crate::FftwPlanner::build_plan(n_fft.get())?;
800            crate::FftwPlan::new(Arc::new(plan))
801        };
802
803        // Apply window and compute FFT
804        let mut windowed = vec![0.0; n_fft.get()];
805        for i in 0..samples.len().get() {
806            windowed[i] = samples[i] * window_samples[i];
807        }
808        // The rest is already zero-padded
809        let mut fft_out = vec![Complex::new(0.0, 0.0); out_len];
810        fft.process(&windowed, &mut fft_out)?;
811
812        // Convert to power
813        let power: Vec<f64> = fft_out.iter().map(num_complex::Complex::norm_sqr).collect();
814        // safety: power is non-empty since n_fft > 0
815        Ok(unsafe { NonEmptyVec::new_unchecked(power) })
816    }
817
818    /// Compute the magnitude spectrum of a single audio frame.
819    ///
820    /// This is useful for real-time processing or analyzing individual frames.
821    ///
822    /// # Arguments
823    ///
824    /// * `samples` - Audio frame (length ≤ n_fft, will be zero-padded if shorter)
825    /// * `n_fft` - FFT size
826    /// * `window` - Window type to apply
827    ///
828    /// # Returns
829    ///
830    /// A vector of magnitude values (|X|) with length `n_fft/2` + 1.
831    ///
832    /// # Automatic Zero-Padding
833    ///
834    /// If the input signal is shorter than `n_fft`, it will be automatically
835    /// zero-padded to the required length.
836    ///
837    /// # Errors
838    ///
839    /// Returns `InvalidInput` error if the input length exceeds `n_fft`.
840    ///
841    /// # Examples
842    ///
843    /// ```
844    /// use spectrograms::*;
845    /// use non_empty_slice::non_empty_vec;
846    ///
847    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
848    /// let frame = non_empty_vec![0.0; nzu!(512)];
849    ///
850    /// let planner = SpectrogramPlanner::new();
851    /// let magnitude = planner.compute_magnitude_spectrum(frame.as_ref(), nzu!(512), WindowType::Hanning)?;
852    ///
853    /// assert_eq!(magnitude.len(), nzu!(257)); // 512/2 + 1
854    /// # Ok(())
855    /// # }
856    /// ```
857    #[inline]
858    pub fn compute_magnitude_spectrum(
859        &self,
860        samples: &NonEmptySlice<f64>,
861        n_fft: NonZeroUsize,
862        window: WindowType,
863    ) -> SpectrogramResult<NonEmptyVec<f64>> {
864        let power = self.compute_power_spectrum(samples, n_fft, window)?;
865        let power = power.iter().map(|&p| p.sqrt()).collect::<Vec<f64>>();
866        // safety: power is non-empty since power_spectrum returned successfully
867        Ok(unsafe { NonEmptyVec::new_unchecked(power) })
868    }
869
870    /// Build a linear-frequency spectrogram plan.
871    ///
872    /// # Type Parameters
873    ///
874    /// `AmpScale` determines whether output is:
875    /// - Magnitude
876    /// - Power
877    /// - Decibels
878    ///
879    /// # Arguments
880    ///
881    /// * `params` - Spectrogram parameters
882    /// * `db` - Logarithmic scaling parameters (only used if `AmpScale
883    /// is `Decibels`)
884    ///
885    /// # Returns
886    ///
887    /// A `SpectrogramPlan` configured for linear-frequency spectrogram computation.
888    ///
889    /// # Errors
890    ///
891    /// Returns an error if the plan cannot be created due to invalid parameters.
892    #[inline]
893    pub fn linear_plan<AmpScale, T: Sample>(
894        &self,
895        params: &SpectrogramParams,
896        db: Option<&LogParams>, // only used when AmpScale = Decibels
897    ) -> SpectrogramResult<SpectrogramPlan<LinearHz, AmpScale, T>>
898    where
899        AmpScale: AmpScaleSpec + 'static,
900    {
901        let stft = StftPlan::<T>::new(params)?;
902        let mapping = FrequencyMapping::<LinearHz>::new(params)?;
903        let scaling = AmplitudeScaling::<AmpScale>::new(db);
904        let freq_axis = build_frequency_axis::<LinearHz>(params, &mapping);
905
906        let workspace = Workspace::<T>::new(stft.n_fft, stft.out_len, mapping.output_bins());
907
908        Ok(SpectrogramPlan {
909            params: params.clone(),
910            stft,
911            mapping,
912            scaling,
913            freq_axis,
914            workspace,
915            _amp: PhantomData,
916        })
917    }
918
919    /// Build a mel-frequency spectrogram plan.
920    ///
921    /// This compiles a mel filterbank matrix and caches it inside the plan.
922    ///
923    /// # Type Parameters
924    ///
925    /// `AmpScale`: determines whether output is:
926    /// - Magnitude
927    /// - Power
928    /// - Decibels
929    ///
930    /// # Arguments
931    ///
932    /// * `params` - Spectrogram parameters
933    /// * `mel` - Mel-specific parameters
934    /// * `db` - Logarithmic scaling parameters (only used if `AmpScale` is `Decibels`)
935    ///
936    /// # Returns
937    ///
938    /// A `SpectrogramPlan` configured for mel spectrogram computation.
939    ///
940    /// # Errors
941    ///
942    /// Returns an error if the plan cannot be created due to invalid parameters.
943    #[inline]
944    pub fn mel_plan<AmpScale, T: Sample>(
945        &self,
946        params: &SpectrogramParams,
947        mel: &MelParams,
948        db: Option<&LogParams>, // only used when AmpScale = Decibels
949    ) -> SpectrogramResult<SpectrogramPlan<Mel, AmpScale, T>>
950    where
951        AmpScale: AmpScaleSpec + 'static,
952    {
953        // cross-validation: mel range must be compatible with sample rate
954        let nyquist = params.nyquist_hz();
955        if mel.f_max() > nyquist {
956            return Err(SpectrogramError::invalid_input(
957                "mel f_max must be <= Nyquist",
958            ));
959        }
960
961        let stft = StftPlan::<T>::new(params)?;
962        let mapping = FrequencyMapping::<Mel>::new_mel(params, mel)?;
963        let scaling = AmplitudeScaling::<AmpScale>::new(db);
964        let freq_axis = build_frequency_axis::<Mel>(params, &mapping);
965
966        let workspace = Workspace::<T>::new(stft.n_fft, stft.out_len, mapping.output_bins());
967
968        Ok(SpectrogramPlan {
969            params: params.clone(),
970            stft,
971            mapping,
972            scaling,
973            freq_axis,
974            workspace,
975            _amp: PhantomData,
976        })
977    }
978
979    /// Build an ERB-scale spectrogram plan.
980    ///
981    /// This creates a spectrogram with ERB-spaced frequency bands using gammatone
982    /// filterbank approximation in the frequency domain.
983    ///
984    /// # Type Parameters
985    ///
986    /// `AmpScale`: determines whether output is:
987    /// - Magnitude
988    /// - Power
989    /// - Decibels
990    ///
991    /// # Arguments
992    ///
993    /// * `params` - Spectrogram parameters
994    /// * `erb` - ERB-specific parameters
995    /// * `db` - Logarithmic scaling parameters (only used if `AmpScale` is `Decibels`)
996    ///
997    /// # Returns
998    ///
999    /// A `SpectrogramPlan` configured for ERB spectrogram computation.
1000    ///
1001    /// # Errors
1002    ///
1003    /// Returns an error if the plan cannot be created due to invalid parameters.
1004    #[inline]
1005    pub fn erb_plan<AmpScale, T: Sample>(
1006        &self,
1007        params: &SpectrogramParams,
1008        erb: &ErbParams,
1009        db: Option<&LogParams>,
1010    ) -> SpectrogramResult<SpectrogramPlan<Erb, AmpScale, T>>
1011    where
1012        AmpScale: AmpScaleSpec + 'static,
1013    {
1014        // cross-validation: erb range must be compatible with sample rate
1015        let nyquist = params.nyquist_hz();
1016        if erb.f_max() > nyquist {
1017            return Err(SpectrogramError::invalid_input(format!(
1018                "f_max={} exceeds Nyquist={}",
1019                erb.f_max(),
1020                nyquist
1021            )));
1022        }
1023
1024        let stft = StftPlan::<T>::new(params)?;
1025        let mapping = FrequencyMapping::<Erb>::new_erb(params, erb)?;
1026        let scaling = AmplitudeScaling::<AmpScale>::new(db);
1027        let freq_axis = build_frequency_axis::<Erb>(params, &mapping);
1028
1029        let workspace = Workspace::<T>::new(stft.n_fft, stft.out_len, mapping.output_bins());
1030
1031        Ok(SpectrogramPlan {
1032            params: params.clone(),
1033            stft,
1034            mapping,
1035            scaling,
1036            freq_axis,
1037            workspace,
1038            _amp: PhantomData,
1039        })
1040    }
1041
1042    /// Build a log-frequency plan.
1043    ///
1044    /// This creates a spectrogram with logarithmically-spaced frequency bins.
1045    ///
1046    /// # Type Parameters
1047    ///
1048    /// `AmpScale`: determines whether output is:
1049    /// - Magnitude
1050    /// - Power
1051    /// - Decibels
1052    ///
1053    /// # Arguments
1054    ///
1055    /// * `params` - Spectrogram parameters
1056    /// * `loghz` - LogHz-specific parameters
1057    /// * `db` - Logarithmic scaling parameters (only used if `AmpScale` is `Decibels`)
1058    ///
1059    /// # Returns
1060    ///
1061    /// A `SpectrogramPlan` configured for log-frequency spectrogram computation.
1062    ///
1063    /// # Errors
1064    ///
1065    /// Returns an error if the plan cannot be created due to invalid parameters.
1066    #[inline]
1067    pub fn log_hz_plan<AmpScale, T: Sample>(
1068        &self,
1069        params: &SpectrogramParams,
1070        loghz: &LogHzParams,
1071        db: Option<&LogParams>,
1072    ) -> SpectrogramResult<SpectrogramPlan<LogHz, AmpScale, T>>
1073    where
1074        AmpScale: AmpScaleSpec + 'static,
1075    {
1076        // cross-validation: loghz range must be compatible with sample rate
1077        let nyquist = params.nyquist_hz();
1078        if loghz.f_max() > nyquist {
1079            return Err(SpectrogramError::invalid_input(format!(
1080                "f_max={} exceeds Nyquist={}",
1081                loghz.f_max(),
1082                nyquist
1083            )));
1084        }
1085
1086        let stft = StftPlan::<T>::new(params)?;
1087        let mapping = FrequencyMapping::<LogHz>::new_loghz(params, loghz)?;
1088        let scaling = AmplitudeScaling::<AmpScale>::new(db);
1089        let freq_axis = build_frequency_axis::<LogHz>(params, &mapping);
1090
1091        let workspace = Workspace::<T>::new(stft.n_fft, stft.out_len, mapping.output_bins());
1092
1093        Ok(SpectrogramPlan {
1094            params: params.clone(),
1095            stft,
1096            mapping,
1097            scaling,
1098            freq_axis,
1099            workspace,
1100            _amp: PhantomData,
1101        })
1102    }
1103
1104    /// Build a cqt spectrogram plan.
1105    ///
1106    /// # Type Parameters
1107    ///
1108    /// `AmpScale`: determines whether output is:
1109    /// - Magnitude
1110    /// - Power
1111    /// - Decibels
1112    ///
1113    /// # Arguments
1114    ///
1115    /// * `params` - Spectrogram parameters
1116    /// * `cqt` - CQT-specific parameters
1117    /// * `db` - Logarithmic scaling parameters (only used if `AmpScale` is `Decibels`)
1118    ///
1119    /// # Returns
1120    ///
1121    /// A `SpectrogramPlan` configured for CQT spectrogram computation.
1122    ///
1123    /// # Errors
1124    ///
1125    /// Returns an error if the plan cannot be created due to invalid parameters.
1126    #[inline]
1127    pub fn cqt_plan<AmpScale, T: Sample>(
1128        &self,
1129        params: &SpectrogramParams,
1130        cqt: &CqtParams,
1131        db: Option<&LogParams>, // only used when AmpScale = Decibels
1132    ) -> SpectrogramResult<SpectrogramPlan<Cqt, AmpScale, T>>
1133    where
1134        AmpScale: AmpScaleSpec + 'static,
1135    {
1136        let stft = StftPlan::<T>::new(params)?;
1137        let mapping = FrequencyMapping::<Cqt>::new(params, cqt)?;
1138        let scaling = AmplitudeScaling::<AmpScale>::new(db);
1139        let freq_axis = build_frequency_axis::<Cqt>(params, &mapping);
1140
1141        let workspace = Workspace::<T>::new(stft.n_fft, stft.out_len, mapping.output_bins());
1142
1143        Ok(SpectrogramPlan {
1144            params: params.clone(),
1145            stft,
1146            mapping,
1147            scaling,
1148            freq_axis,
1149            workspace,
1150            _amp: PhantomData,
1151        })
1152    }
1153}
1154
1155/// STFT plan containing reusable FFT plan and buffers.
1156///
1157/// This struct is responsible for performing the Short-Time Fourier Transform (STFT)
1158/// on audio signals based on the provided parameters.
1159///
1160/// It encapsulates the FFT plan, windowing function, and internal buffers to efficiently
1161/// compute the STFT for multiple frames of audio data.
1162///
1163/// # Fields
1164///
1165/// - `n_fft`: Size of the FFT.
1166/// - `hop_size`: Hop size between consecutive frames.
1167/// - `window`: Windowing function samples.
1168/// - `centre`: Whether to centre the frames with padding.
1169/// - `out_len`: Length of the FFT output.
1170/// - `fft`: Boxed FFT plan for real-to-complex transformation.
1171/// - `fft_out`: Internal buffer for FFT output.
1172/// - `frame`: Internal buffer for windowed audio frames.
1173pub struct StftPlan<T = f64> {
1174    n_fft: NonZeroUsize,
1175    hop_size: NonZeroUsize,
1176    window: NonEmptyVec<T>,
1177    centre: bool,
1178
1179    out_len: NonZeroUsize,
1180
1181    // FFT plan (reused for all frames)
1182    fft: Box<dyn R2cPlan<T>>,
1183
1184    // internal scratch
1185    fft_out: NonEmptyVec<Complex<T>>,
1186    frame: NonEmptyVec<T>,
1187}
1188
1189impl<T: Sample> StftPlan<T> {
1190    /// Create a new STFT plan from parameters.
1191    ///
1192    /// # Arguments
1193    ///
1194    /// * `params` - Spectrogram parameters containing STFT config
1195    ///
1196    /// # Returns
1197    ///
1198    /// A new `StftPlan` instance.
1199    ///
1200    /// # Errors
1201    ///
1202    /// Returns an error if the FFT plan cannot be created.
1203    #[inline]
1204    pub fn new(params: &SpectrogramParams) -> SpectrogramResult<Self> {
1205        let stft = params.stft();
1206        let n_fft = stft.n_fft();
1207        let hop_size = stft.hop_size();
1208        let centre = stft.centre();
1209
1210        let window = make_window(stft.window(), n_fft);
1211
1212        let out_len = r2c_output_size(n_fft.get());
1213        let out_len = NonZeroUsize::new(out_len)
1214            .ok_or_else(|| SpectrogramError::invalid_input("FFT output length must be non-zero"))?;
1215
1216        let fft: Box<dyn R2cPlan<T>> = Box::new(T::plan_r2c(n_fft.get())?);
1217
1218        Ok(Self {
1219            n_fft,
1220            hop_size,
1221            window,
1222            centre,
1223            out_len,
1224            fft,
1225            fft_out: non_empty_vec![Complex::new(T::zero(), T::zero()); out_len],
1226            frame: non_empty_vec![T::zero(); n_fft],
1227        })
1228    }
1229
1230    fn frame_count(&self, n_samples: NonZeroUsize) -> SpectrogramResult<NonZeroUsize> {
1231        // Framing policy:
1232        // - centre = true: implicit padding of n_fft/2 on both sides
1233        // - centre = false: no padding
1234        //
1235        // Define the number of frames such that each frame has a valid centre sample position.
1236        let pad = if self.centre { self.n_fft.get() / 2 } else { 0 };
1237        let padded_len = n_samples.get() + 2 * pad;
1238
1239        if padded_len < self.n_fft.get() {
1240            // still produce 1 frame (all padding / partial)
1241            return Ok(nzu!(1));
1242        }
1243
1244        let remaining = padded_len - self.n_fft.get();
1245        let n_frames = remaining / self.hop_size().get() + 1;
1246        let n_frames = NonZeroUsize::new(n_frames).ok_or_else(|| {
1247            SpectrogramError::invalid_input("computed number of frames must be non-zero")
1248        })?;
1249        Ok(n_frames)
1250    }
1251
1252    /// Compute one frame FFT using internal buffers only.
1253    fn compute_frame_fft_simple(
1254        &mut self,
1255        samples: &NonEmptySlice<T>,
1256        frame_idx: usize,
1257    ) -> SpectrogramResult<()> {
1258        let out = self.frame.as_mut_slice();
1259        debug_assert_eq!(out.len(), self.n_fft.get());
1260
1261        let pad = if self.centre { self.n_fft.get() / 2 } else { 0 };
1262        let start = frame_idx
1263            .checked_mul(self.hop_size.get())
1264            .ok_or_else(|| SpectrogramError::invalid_input("frame index overflow"))?;
1265
1266        // Fill windowed frame
1267        for (i, sample) in out.iter_mut().enumerate().take(self.n_fft.get()) {
1268            let v_idx = start + i;
1269            let s_idx = v_idx as isize - pad as isize;
1270
1271            let sample_val = if s_idx < 0 || (s_idx as usize) >= samples.len().get() {
1272                T::zero()
1273            } else {
1274                samples[s_idx as usize]
1275            };
1276            *sample = sample_val * self.window[i];
1277        }
1278
1279        // Compute FFT
1280        let fft_out = self.fft_out.as_mut_slice();
1281        self.fft.process(out, fft_out)?;
1282
1283        Ok(())
1284    }
1285
1286    /// Compute one frame spectrum into workspace:
1287    /// - fills windowed frame
1288    /// - runs FFT
1289    /// - converts to magnitude/power based on `AmpScale` later
1290    fn compute_frame_spectrum(
1291        &mut self,
1292        samples: &NonEmptySlice<T>,
1293        frame_idx: usize,
1294        workspace: &mut Workspace<T>,
1295    ) -> SpectrogramResult<()> {
1296        let out = workspace.frame.as_mut_slice();
1297
1298        // self.fill_frame(samples, frame_idx, frame)?;
1299        debug_assert_eq!(out.len(), self.n_fft.get());
1300
1301        let pad = if self.centre { self.n_fft.get() / 2 } else { 0 };
1302        let start = frame_idx
1303            .checked_mul(self.hop_size().get())
1304            .ok_or_else(|| SpectrogramError::invalid_input("frame index overflow"))?;
1305
1306        // The "virtual" signal is samples with pad zeros on both sides.
1307        // Virtual index 0..padded_len
1308        // Map virtual index to original samples by subtracting pad.
1309        for (i, sample) in out.iter_mut().enumerate().take(self.n_fft.get()) {
1310            let v_idx = start + i;
1311            let s_idx = v_idx as isize - pad as isize;
1312
1313            let sample_val = if s_idx < 0 || (s_idx as usize) >= samples.len().get() {
1314                T::zero()
1315            } else {
1316                samples[s_idx as usize]
1317            };
1318
1319            *sample = sample_val * self.window[i];
1320        }
1321        let fft_out = workspace.fft_out.as_mut_slice();
1322        // FFT
1323        self.fft.process(out, fft_out)?;
1324
1325        // Convert complex spectrum to linear magnitude OR power here? No:
1326        // Keep "spectrum" as power by default? That would entangle semantics.
1327        //
1328        // Instead, we store magnitude^2 (power) as the canonical intermediate,
1329        // and let AmpScale decide later whether output is magnitude or power.
1330        //
1331        // This is consistent and avoids recomputing norms multiple times.
1332        for (i, c) in workspace.fft_out.iter().enumerate() {
1333            workspace.spectrum[i] = c.norm_sqr();
1334        }
1335
1336        Ok(())
1337    }
1338
1339    /// Fill a time-domain frame WITHOUT applying the window.
1340    ///
1341    /// This is used for CQT mapping, where the CQT kernels already contain
1342    /// windowing applied during kernel generation. Applying the STFT window
1343    /// would result in double-windowing.
1344    ///
1345    /// # Arguments
1346    ///
1347    /// * `samples` - Input audio samples
1348    /// * `frame_idx` - Frame index
1349    /// * `workspace` - Workspace containing the frame buffer to fill
1350    ///
1351    /// # Errors
1352    ///
1353    /// Returns an error if frame index is out of bounds.
1354    fn fill_frame_unwindowed(
1355        &self,
1356        samples: &NonEmptySlice<T>,
1357        frame_idx: usize,
1358        workspace: &mut Workspace<T>,
1359    ) -> SpectrogramResult<()> {
1360        let out = workspace.frame.as_mut_slice();
1361        debug_assert_eq!(out.len(), self.n_fft.get());
1362
1363        let pad = if self.centre { self.n_fft.get() / 2 } else { 0 };
1364        let start = frame_idx
1365            .checked_mul(self.hop_size().get())
1366            .ok_or_else(|| SpectrogramError::invalid_input("frame index overflow"))?;
1367
1368        // Fill frame WITHOUT windowing
1369        for (i, sample) in out.iter_mut().enumerate().take(self.n_fft.get()) {
1370            let v_idx = start + i;
1371            let s_idx = v_idx as isize - pad as isize;
1372
1373            let sample_val = if s_idx < 0 || (s_idx as usize) >= samples.len().get() {
1374                T::zero()
1375            } else {
1376                samples[s_idx as usize]
1377            };
1378
1379            // No window multiplication - just copy the sample
1380            *sample = sample_val;
1381        }
1382
1383        Ok(())
1384    }
1385
1386    /// Compute the full STFT for a signal, returning an `StftResult`.
1387    ///
1388    /// This is a convenience method that handles frame iteration and
1389    /// builds the complete STFT matrix.
1390    ///
1391    ///
1392    /// # Arguments
1393    ///
1394    /// * `samples` - Input audio samples
1395    /// * `params` - STFT computation parameters
1396    ///
1397    /// # Returns
1398    ///
1399    /// An `StftResult` containing the complex STFT matrix and metadata.
1400    ///
1401    /// # Errors
1402    ///
1403    /// Returns an error if computation fails.
1404    ///
1405    /// # Examples
1406    ///
1407    /// ```rust
1408    /// use spectrograms::*;
1409    /// use non_empty_slice::non_empty_vec;
1410    ///
1411    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1412    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
1413    /// let params = SpectrogramParams::new(stft, 16000.0)?;
1414    /// let mut plan = StftPlan::new(&params)?;
1415    ///
1416    /// let samples = non_empty_vec![0.0; nzu!(16000)];
1417    /// let stft_result = plan.compute(&samples, &params)?;
1418    ///
1419    /// println!("STFT: {} bins x {} frames", stft_result.n_bins(), stft_result.n_frames());
1420    /// # Ok(())
1421    /// # }
1422    /// ```
1423    #[inline]
1424    pub fn compute(
1425        &mut self,
1426        samples: &NonEmptySlice<T>,
1427        params: &SpectrogramParams,
1428    ) -> SpectrogramResult<StftResult<T>> {
1429        let n_frames = self.frame_count(samples.len())?;
1430        let n_bins = self.out_len;
1431
1432        // Allocate output matrix (frequency_bins x time_frames)
1433        let mut data = Array2::<Complex<T>>::zeros((n_bins.get(), n_frames.get()));
1434
1435        // Compute each frame
1436        for frame_idx in 0..n_frames.get() {
1437            self.compute_frame_fft_simple(samples, frame_idx)?;
1438
1439            // Copy from internal buffer to output
1440            for (bin_idx, &value) in self.fft_out.iter().enumerate() {
1441                data[[bin_idx, frame_idx]] = value;
1442            }
1443        }
1444
1445        // Build frequency axis
1446        let frequencies: Vec<f64> = (0..n_bins.get())
1447            .map(|k| k as f64 * params.sample_rate_hz() / params.stft().n_fft().get() as f64)
1448            .collect();
1449        // SAFETY: n_bins > 0
1450        let frequencies = unsafe { NonEmptyVec::new_unchecked(frequencies) };
1451
1452        Ok(StftResult {
1453            data,
1454            frequencies,
1455            sample_rate: params.sample_rate_hz(),
1456            params: params.stft().clone(),
1457        })
1458    }
1459
1460    /// Compute a single frame of STFT, returning the complex spectrum.
1461    ///
1462    /// This is useful for streaming/online processing where you want to
1463    /// process audio frame-by-frame.
1464    ///
1465    /// # Arguments
1466    ///
1467    /// * `samples` - Input audio samples
1468    /// * `frame_idx` - Index of the frame to compute
1469    ///
1470    /// # Returns
1471    ///
1472    /// A `NonEmptyVec` containing the complex spectrum for the specified frame.
1473    ///
1474    /// # Errors
1475    ///
1476    /// Returns an error if computation fails.
1477    ///
1478    /// # Examples
1479    ///
1480    /// ```rust
1481    /// use spectrograms::*;
1482    /// use non_empty_slice::non_empty_vec;
1483    ///
1484    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1485    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
1486    /// let params = SpectrogramParams::new(stft, 16000.0)?;
1487    /// let mut plan = StftPlan::new(&params)?;
1488    ///
1489    /// let samples = non_empty_vec![0.0; nzu!(16000)];
1490    /// let (_, n_frames) = plan.output_shape(samples.len())?;
1491    ///
1492    /// for frame_idx in 0..n_frames.get() {
1493    ///     let spectrum = plan.compute_frame_simple(&samples, frame_idx)?;
1494    ///     // Process spectrum...
1495    /// }
1496    /// # Ok(())
1497    /// # }
1498    /// ```
1499    #[inline]
1500    pub fn compute_frame_simple(
1501        &mut self,
1502        samples: &NonEmptySlice<T>,
1503        frame_idx: usize,
1504    ) -> SpectrogramResult<NonEmptyVec<Complex<T>>> {
1505        self.compute_frame_fft_simple(samples, frame_idx)?;
1506        Ok(self.fft_out.clone())
1507    }
1508
1509    /// Compute STFT into a pre-allocated buffer.
1510    ///
1511    /// This avoids allocating the output matrix, useful for reusing buffers.
1512    ///
1513    /// # Arguments
1514    ///
1515    /// * `samples` - Input audio samples
1516    /// * `output` - Pre-allocated output buffer (shape: `n_bins` x `n_frames`)
1517    ///  
1518    /// # Returns
1519    ///
1520    /// `Ok(())` on success, or an error if dimensions mismatch.
1521    ///
1522    /// # Errors
1523    ///
1524    /// Returns `DimensionMismatch` error if the output buffer has incorrect shape.
1525    ///
1526    /// # Examples
1527    ///
1528    /// ```rust
1529    /// use spectrograms::*;
1530    /// use ndarray::Array2;
1531    /// use num_complex::Complex;
1532    /// use non_empty_slice::non_empty_vec;
1533    ///
1534    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1535    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
1536    /// let params = SpectrogramParams::new(stft, 16000.0)?;
1537    /// let mut plan = StftPlan::new(&params)?;
1538    ///
1539    /// let samples = non_empty_vec![0.0; nzu!(16000)];
1540    /// let (n_bins, n_frames) = plan.output_shape(samples.len())?;
1541    /// let mut output = Array2::<Complex<f64>>::zeros((n_bins.get(), n_frames.get()));
1542    ///
1543    /// plan.compute_into(&samples, &mut output)?;
1544    /// # Ok(())
1545    /// # }
1546    /// ```
1547    #[inline]
1548    pub fn compute_into(
1549        &mut self,
1550        samples: &NonEmptySlice<T>,
1551        output: &mut Array2<Complex<T>>,
1552    ) -> SpectrogramResult<()> {
1553        let n_frames = self.frame_count(samples.len())?;
1554        let n_bins = self.out_len;
1555
1556        // Validate output dimensions
1557        if output.nrows() != n_bins.get() {
1558            return Err(SpectrogramError::dimension_mismatch(
1559                n_bins.get(),
1560                output.nrows(),
1561            ));
1562        }
1563        if output.ncols() != n_frames.get() {
1564            return Err(SpectrogramError::dimension_mismatch(
1565                n_frames.get(),
1566                output.ncols(),
1567            ));
1568        }
1569
1570        // Compute into pre-allocated buffer
1571        for frame_idx in 0..n_frames.get() {
1572            self.compute_frame_fft_simple(samples, frame_idx)?;
1573
1574            for (bin_idx, &value) in self.fft_out.iter().enumerate() {
1575                output[[bin_idx, frame_idx]] = value;
1576            }
1577        }
1578
1579        Ok(())
1580    }
1581
1582    /// Get the expected output dimensions for a given signal length.
1583    ///
1584    /// # Arguments
1585    ///
1586    /// * `signal_length` - Length of the input signal in samples
1587    ///
1588    /// # Returns
1589    ///
1590    /// A tuple `(n_frequency_bins, n_time_frames)`.
1591    ///
1592    /// # Errors
1593    ///
1594    /// Returns an error if the computed number of frames is invalid.
1595    #[inline]
1596    pub fn output_shape(
1597        &self,
1598        signal_length: NonZeroUsize,
1599    ) -> SpectrogramResult<(NonZeroUsize, NonZeroUsize)> {
1600        let n_frames = self.frame_count(signal_length)?;
1601        Ok((self.out_len, n_frames))
1602    }
1603
1604    /// Get the number of frequency bins in the output.
1605    ///
1606    /// # Returns
1607    ///
1608    /// The number of frequency bins.
1609    #[inline]
1610    #[must_use]
1611    pub const fn n_bins(&self) -> NonZeroUsize {
1612        self.out_len
1613    }
1614
1615    /// Get the FFT size.
1616    ///
1617    /// # Returns
1618    ///
1619    /// The FFT size.
1620    #[inline]
1621    #[must_use]
1622    pub const fn n_fft(&self) -> NonZeroUsize {
1623        self.n_fft
1624    }
1625
1626    /// Get the hop size.
1627    ///
1628    /// # Returns
1629    ///
1630    /// The hop size.
1631    #[inline]
1632    #[must_use]
1633    pub const fn hop_size(&self) -> NonZeroUsize {
1634        self.hop_size
1635    }
1636}
1637
1638#[derive(Debug, Clone)]
1639enum MappingKind {
1640    Identity {
1641        out_len: NonZeroUsize,
1642    },
1643    Mel {
1644        matrix: SparseMatrix,
1645    }, // shape: (n_mels, out_len)
1646    LogHz {
1647        matrix: SparseMatrix,
1648        frequencies: NonEmptyVec<f64>,
1649    }, // shape: (n_bins, out_len)
1650    Erb {
1651        filterbank: ErbFilterbank,
1652    },
1653    Cqt {
1654        kernel: CqtKernel,
1655    },
1656}
1657
1658impl MappingKind {
1659    /// Check if this mapping requires unwindowed time-domain frames.
1660    ///
1661    /// CQT kernels already contain windowing applied during kernel generation,
1662    /// so they should receive unwindowed frames to avoid double-windowing.
1663    /// All other mappings work on FFT spectra and don't use the time-domain frame.
1664    const fn needs_unwindowed_frame(&self) -> bool {
1665        matches!(self, Self::Cqt { .. })
1666    }
1667}
1668
1669/// Typed mapping wrapper.
1670#[derive(Debug, Clone)]
1671struct FrequencyMapping<FreqScale> {
1672    kind: MappingKind,
1673    _marker: PhantomData<FreqScale>,
1674}
1675
1676impl FrequencyMapping<LinearHz> {
1677    fn new(params: &SpectrogramParams) -> SpectrogramResult<Self> {
1678        let out_len = r2c_output_size(params.stft().n_fft().get());
1679        let out_len = NonZeroUsize::new(out_len)
1680            .ok_or_else(|| SpectrogramError::invalid_input("FFT output length must be non-zero"))?;
1681        Ok(Self {
1682            kind: MappingKind::Identity { out_len },
1683            _marker: PhantomData,
1684        })
1685    }
1686}
1687
1688impl FrequencyMapping<Mel> {
1689    fn new_mel(params: &SpectrogramParams, mel: &MelParams) -> SpectrogramResult<Self> {
1690        let n_fft = params.stft().n_fft();
1691        let out_len = r2c_output_size(n_fft.get());
1692        let out_len = NonZeroUsize::new(out_len)
1693            .ok_or_else(|| SpectrogramError::invalid_input("FFT output length must be non-zero"))?;
1694
1695        // Validate: mel bins must be <= something sensible
1696        if mel.n_mels() > nzu!(10_000) {
1697            return Err(SpectrogramError::invalid_input(
1698                "n_mels is unreasonably large",
1699            ));
1700        }
1701
1702        let matrix = build_mel_filterbank_matrix(
1703            params.sample_rate_hz(),
1704            n_fft,
1705            mel.n_mels(),
1706            mel.f_min(),
1707            mel.f_max(),
1708            mel.norm(),
1709        )?;
1710
1711        // matrix must be (n_mels, out_len)
1712        if matrix.nrows() != mel.n_mels().get() || matrix.ncols() != out_len.get() {
1713            return Err(SpectrogramError::invalid_input(
1714                "mel filterbank matrix shape mismatch",
1715            ));
1716        }
1717
1718        Ok(Self {
1719            kind: MappingKind::Mel { matrix },
1720            _marker: PhantomData,
1721        })
1722    }
1723}
1724
1725impl FrequencyMapping<LogHz> {
1726    fn new_loghz(params: &SpectrogramParams, loghz: &LogHzParams) -> SpectrogramResult<Self> {
1727        let n_fft = params.stft().n_fft();
1728        let out_len = r2c_output_size(n_fft.get());
1729        let out_len = NonZeroUsize::new(out_len)
1730            .ok_or_else(|| SpectrogramError::invalid_input("FFT output length must be non-zero"))?;
1731        // Validate: n_bins must be <= something sensible
1732        if loghz.n_bins() > nzu!(10_000) {
1733            return Err(SpectrogramError::invalid_input(
1734                "n_bins is unreasonably large",
1735            ));
1736        }
1737
1738        let (matrix, frequencies) = build_loghz_matrix(
1739            params.sample_rate_hz(),
1740            n_fft,
1741            loghz.n_bins(),
1742            loghz.f_min(),
1743            loghz.f_max(),
1744        )?;
1745
1746        // matrix must be (n_bins, out_len)
1747        if matrix.nrows() != loghz.n_bins().get() || matrix.ncols() != out_len.get() {
1748            return Err(SpectrogramError::invalid_input(
1749                "loghz matrix shape mismatch",
1750            ));
1751        }
1752
1753        Ok(Self {
1754            kind: MappingKind::LogHz {
1755                matrix,
1756                frequencies,
1757            },
1758            _marker: PhantomData,
1759        })
1760    }
1761}
1762
1763impl FrequencyMapping<Erb> {
1764    fn new_erb(params: &SpectrogramParams, erb: &crate::erb::ErbParams) -> SpectrogramResult<Self> {
1765        let n_fft = params.stft().n_fft();
1766        let sample_rate = params.sample_rate_hz();
1767
1768        // Validate: n_filters must be <= something sensible
1769        if erb.n_filters() > nzu!(10_000) {
1770            return Err(SpectrogramError::invalid_input(
1771                "n_filters is unreasonably large",
1772            ));
1773        }
1774
1775        // Generate ERB filterbank with pre-computed frequency responses
1776        let filterbank = crate::erb::ErbFilterbank::generate(erb, sample_rate, n_fft)?;
1777
1778        Ok(Self {
1779            kind: MappingKind::Erb { filterbank },
1780            _marker: PhantomData,
1781        })
1782    }
1783}
1784
1785impl FrequencyMapping<Cqt> {
1786    fn new(params: &SpectrogramParams, cqt: &CqtParams) -> SpectrogramResult<Self> {
1787        let sample_rate = params.sample_rate_hz();
1788        let n_fft = params.stft().n_fft();
1789
1790        // Validate that frequency range is reasonable
1791        let f_max = cqt.bin_frequency(cqt.num_bins().get().saturating_sub(1));
1792        if f_max >= sample_rate / 2.0 {
1793            return Err(SpectrogramError::invalid_input(
1794                "CQT maximum frequency must be below Nyquist frequency",
1795            ));
1796        }
1797
1798        // Generate CQT kernel using n_fft as the signal length for kernel generation
1799        let kernel = CqtKernel::generate(cqt, sample_rate, n_fft);
1800
1801        Ok(Self {
1802            kind: MappingKind::Cqt { kernel },
1803            _marker: PhantomData,
1804        })
1805    }
1806}
1807
1808impl<FreqScale> FrequencyMapping<FreqScale> {
1809    const fn output_bins(&self) -> NonZeroUsize {
1810        // safety: all variants ensure output bins > 0 OR rely on a matrix that is guaranteed to have rows > 0
1811        match &self.kind {
1812            MappingKind::Identity { out_len } => *out_len,
1813            // safety: matrix.nrows() > 0
1814            MappingKind::LogHz { matrix, .. } | MappingKind::Mel { matrix } => unsafe {
1815                NonZeroUsize::new_unchecked(matrix.nrows())
1816            },
1817            MappingKind::Erb { filterbank, .. } => filterbank.num_filters(),
1818            MappingKind::Cqt { kernel, .. } => kernel.num_bins(),
1819        }
1820    }
1821
1822    fn apply<T: Sample>(
1823        &self,
1824        spectrum: &NonEmptySlice<T>,
1825        frame: &NonEmptySlice<T>,
1826        out: &mut NonEmptySlice<T>,
1827    ) -> SpectrogramResult<()> {
1828        match &self.kind {
1829            MappingKind::Identity { out_len } => {
1830                if spectrum.len() != *out_len {
1831                    return Err(SpectrogramError::dimension_mismatch(
1832                        (*out_len).get(),
1833                        spectrum.len().get(),
1834                    ));
1835                }
1836                if out.len() != *out_len {
1837                    return Err(SpectrogramError::dimension_mismatch(
1838                        (*out_len).get(),
1839                        out.len().get(),
1840                    ));
1841                }
1842                out.copy_from_slice(spectrum);
1843                Ok(())
1844            }
1845            MappingKind::LogHz { matrix, .. } | MappingKind::Mel { matrix } => {
1846                let out_bins = matrix.nrows();
1847                let in_bins = matrix.ncols();
1848
1849                if spectrum.len().get() != in_bins {
1850                    return Err(SpectrogramError::dimension_mismatch(
1851                        in_bins,
1852                        spectrum.len().get(),
1853                    ));
1854                }
1855                if out.len().get() != out_bins {
1856                    return Err(SpectrogramError::dimension_mismatch(
1857                        out_bins,
1858                        out.len().get(),
1859                    ));
1860                }
1861
1862                // Sparse matrix-vector multiplication: out = matrix * spectrum
1863                matrix.multiply_vec(spectrum, out);
1864                Ok(())
1865            }
1866            MappingKind::Erb { filterbank } => {
1867                // Apply ERB filterbank using pre-computed frequency responses
1868                // The filterbank already has |H(f)|^2 pre-computed, so we just
1869                // apply it to the power spectrum
1870                let erb_out = filterbank.apply_to_power_spectrum(spectrum)?;
1871
1872                if out.len().get() != erb_out.len().get() {
1873                    return Err(SpectrogramError::dimension_mismatch(
1874                        erb_out.len().get(),
1875                        out.len().get(),
1876                    ));
1877                }
1878
1879                out.copy_from_slice(&erb_out);
1880                Ok(())
1881            }
1882            MappingKind::Cqt { kernel } => {
1883                // CQT works on time-domain unwindowed frame (not FFT spectrum).
1884                // The CQT kernels contain windowing applied during kernel generation,
1885                // so the input frame should be unwindowed to avoid double-windowing.
1886                // The kernel coefficients are stored in f64 and converted to `T` at
1887                // apply time inside `CqtKernel::apply`, which already operates on a
1888                // `T` frame and yields `Complex<T>` coefficients.
1889                let cqt_complex = kernel.apply(frame)?;
1890
1891                if out.len().get() != cqt_complex.len().get() {
1892                    return Err(SpectrogramError::dimension_mismatch(
1893                        cqt_complex.len().get(),
1894                        out.len().get(),
1895                    ));
1896                }
1897
1898                // Convert complex coefficients to power (|z|^2)
1899                // This matches the convention where intermediate values are in power domain
1900                for (i, c) in cqt_complex.iter().enumerate() {
1901                    out[i] = c.norm_sqr();
1902                }
1903
1904                Ok(())
1905            }
1906        }
1907    }
1908
1909    fn frequencies_hz(&self, params: &SpectrogramParams) -> NonEmptyVec<f64> {
1910        match &self.kind {
1911            MappingKind::Identity { out_len } => {
1912                // Standard R2C bins: k * sr / n_fft
1913                let n_fft = params.stft().n_fft().get() as f64;
1914                let sr = params.sample_rate_hz();
1915                let df = sr / n_fft;
1916
1917                let mut f = Vec::with_capacity((*out_len).get());
1918                for k in 0..(*out_len).get() {
1919                    f.push(k as f64 * df);
1920                }
1921                // safety: out_len > 0
1922                unsafe { NonEmptyVec::new_unchecked(f) }
1923            }
1924            MappingKind::Mel { matrix } => {
1925                // For mel, the axis is defined by the mel band centre frequencies.
1926                // We compute and store them consistently with how we built the filterbank.
1927                let n_mels = matrix.nrows();
1928                // safety: n_mels > 0
1929                let n_mels = unsafe { NonZeroUsize::new_unchecked(n_mels) };
1930                mel_band_centres_hz(n_mels, params.sample_rate_hz(), params.nyquist_hz())
1931            }
1932            MappingKind::LogHz { frequencies, .. } => {
1933                // Frequencies are stored when the mapping is created
1934                frequencies.clone()
1935            }
1936            MappingKind::Erb { filterbank, .. } => {
1937                // ERB center frequencies
1938                filterbank.center_frequencies().to_non_empty_vec()
1939            }
1940            MappingKind::Cqt { kernel, .. } => {
1941                // CQT center frequencies from the kernel
1942                kernel.frequencies().to_non_empty_vec()
1943            }
1944        }
1945    }
1946}
1947
1948//
1949// ========================
1950// Amplitude scaling
1951// ========================
1952//
1953
1954/// Marker trait so we can specialise behaviour by `AmpScale`.
1955pub trait AmpScaleSpec: Sized + Send + Sync {
1956    /// Apply conversion from power-domain value to the desired amplitude scale.
1957    ///
1958    /// # Arguments
1959    ///
1960    /// - `power`: input power-domain value.
1961    ///
1962    /// # Returns
1963    ///
1964    /// Converted amplitude value.
1965    fn apply_from_power<T: Sample>(power: T) -> T;
1966
1967    /// Apply dB conversion in-place on a power-domain vector.
1968    ///
1969    /// This is a no-op for Power and Magnitude scales.
1970    ///
1971    /// # Arguments
1972    ///
1973    /// - `x`: power-domain values to convert to dB in-place.
1974    /// - `floor_db`: dB floor value to apply.
1975    ///
1976    /// # Returns
1977    ///
1978    /// `SpectrogramResult<()>`: Ok on success, error on invalid input.
1979    ///
1980    /// # Errors
1981    ///
1982    /// - If `floor_db` is not finite.
1983    fn apply_db_in_place<T: Sample>(x: &mut [T], floor_db: f64) -> SpectrogramResult<()>;
1984}
1985
1986impl AmpScaleSpec for Power {
1987    #[inline]
1988    fn apply_from_power<T: Sample>(power: T) -> T {
1989        power
1990    }
1991
1992    #[inline]
1993    fn apply_db_in_place<T: Sample>(_x: &mut [T], _floor_db: f64) -> SpectrogramResult<()> {
1994        Ok(())
1995    }
1996}
1997
1998impl AmpScaleSpec for Magnitude {
1999    #[inline]
2000    fn apply_from_power<T: Sample>(power: T) -> T {
2001        power.sqrt()
2002    }
2003
2004    #[inline]
2005    fn apply_db_in_place<T: Sample>(_x: &mut [T], _floor_db: f64) -> SpectrogramResult<()> {
2006        Ok(())
2007    }
2008}
2009
2010impl AmpScaleSpec for Decibels {
2011    #[inline]
2012    fn apply_from_power<T: Sample>(power: T) -> T {
2013        // dB conversion is applied in batch, not here.
2014        power
2015    }
2016
2017    #[inline]
2018    fn apply_db_in_place<T: Sample>(x: &mut [T], floor_db: f64) -> SpectrogramResult<()> {
2019        // Convert power -> dB: 10*log10(max(power, eps))
2020        // where eps is derived from floor_db to ensure consistency
2021        if !floor_db.is_finite() {
2022            return Err(SpectrogramError::invalid_input("floor_db must be finite"));
2023        }
2024
2025        // Convert floor_db to linear scale to get epsilon
2026        // e.g., floor_db = -80 dB -> eps = 10^(-80/10) = 1e-8
2027        // The dB constants are evaluated in f64 (unchanged) and converted to `T`.
2028        let eps = T::from_f64(10.0_f64.powf(floor_db / 10.0));
2029        let ten = T::from_f64(10.0);
2030
2031        for v in x.iter_mut() {
2032            // Clamp power to epsilon before log to avoid log(0) and ensure floor
2033            *v = ten * v.max(eps).log10();
2034        }
2035        Ok(())
2036    }
2037}
2038
2039/// Amplitude scaling configuration.
2040///
2041/// This handles conversion from power-domain intermediate to the desired amplitude scale (Power, Magnitude, Decibels).
2042#[derive(Debug, Clone)]
2043struct AmplitudeScaling<AmpScale> {
2044    db_floor: Option<f64>,
2045    _marker: PhantomData<AmpScale>,
2046}
2047
2048impl<AmpScale> AmplitudeScaling<AmpScale>
2049where
2050    AmpScale: AmpScaleSpec + 'static,
2051{
2052    fn new(db: Option<&LogParams>) -> Self {
2053        let db_floor = db.map(LogParams::floor_db);
2054        Self {
2055            db_floor,
2056            _marker: PhantomData,
2057        }
2058    }
2059
2060    /// Apply amplitude scaling in-place on a mapped spectrum vector.
2061    ///
2062    /// The input vector is assumed to be in the *power* domain (|X|^2),
2063    /// because the STFT stage produces power as the canonical intermediate.
2064    ///
2065    /// - Power: leaves values unchanged.
2066    /// - Magnitude: sqrt(power).
2067    /// - Decibels: converts power -> dB and floors at `db_floor`.
2068    pub fn apply_in_place<T: Sample>(&self, x: &mut [T]) -> SpectrogramResult<()> {
2069        // Convert from canonical power-domain intermediate into the requested linear domain.
2070        for v in x.iter_mut() {
2071            *v = AmpScale::apply_from_power(*v);
2072        }
2073
2074        // Apply dB conversion if configured (no-op for Power/Magnitude via trait impls).
2075        if let Some(floor_db) = self.db_floor {
2076            AmpScale::apply_db_in_place(x, floor_db)?;
2077        }
2078
2079        Ok(())
2080    }
2081}
2082
2083#[derive(Debug, Clone)]
2084struct Workspace<T = f64> {
2085    spectrum: NonEmptyVec<T>,         // out_len (power spectrum, native T)
2086    mapped: NonEmptyVec<T>,           // n_bins (after mapping, native T)
2087    frame: NonEmptyVec<T>,            // n_fft (windowed frame for FFT)
2088    fft_out: NonEmptyVec<Complex<T>>, // out_len (FFT output)
2089}
2090
2091impl<T: Sample> Workspace<T> {
2092    fn new(n_fft: NonZeroUsize, out_len: NonZeroUsize, n_bins: NonZeroUsize) -> Self {
2093        Self {
2094            spectrum: non_empty_vec![T::zero(); out_len],
2095            mapped: non_empty_vec![T::zero(); n_bins],
2096            frame: non_empty_vec![T::zero(); n_fft],
2097            fft_out: non_empty_vec![Complex::new(T::zero(), T::zero()); out_len],
2098        }
2099    }
2100
2101    fn ensure_sizes(&mut self, n_fft: NonZeroUsize, out_len: NonZeroUsize, n_bins: NonZeroUsize) {
2102        if self.spectrum.len() != out_len {
2103            self.spectrum.resize(out_len, T::zero());
2104        }
2105        if self.mapped.len() != n_bins {
2106            self.mapped.resize(n_bins, T::zero());
2107        }
2108        if self.frame.len() != n_fft {
2109            self.frame.resize(n_fft, T::zero());
2110        }
2111        if self.fft_out.len() != out_len {
2112            self.fft_out
2113                .resize(out_len, Complex::new(T::zero(), T::zero()));
2114        }
2115    }
2116}
2117
2118fn build_frequency_axis<FreqScale>(
2119    params: &SpectrogramParams,
2120    mapping: &FrequencyMapping<FreqScale>,
2121) -> FrequencyAxis<FreqScale>
2122where
2123    FreqScale: Copy + Clone + 'static,
2124{
2125    let frequencies = mapping.frequencies_hz(params);
2126    FrequencyAxis::new(frequencies)
2127}
2128
2129fn build_time_axis_seconds(params: &SpectrogramParams, n_frames: NonZeroUsize) -> NonEmptyVec<f64> {
2130    let dt = params.frame_period_seconds();
2131    let mut times = Vec::with_capacity(n_frames.get());
2132
2133    for i in 0..n_frames.get() {
2134        times.push(i as f64 * dt);
2135    }
2136
2137    // safety: times is guaranteed non-empty since n_frames > 0
2138
2139    unsafe { NonEmptyVec::new_unchecked(times) }
2140}
2141
2142/// Generate window function samples.
2143///
2144/// Supports various window types including Rectangular, Hanning, Hamming, Blackman, Kaiser, and Gaussian.
2145///
2146/// # Arguments
2147///
2148/// * `window` - The type of window function to generate.
2149/// * `n_fft` - The size of the FFT, which determines the length of the window.
2150///
2151/// # Returns
2152///
2153/// A `NonEmptyVec<f64>` containing the window function samples.
2154///
2155/// # Panics
2156///
2157/// Panics if a custom window is provided with a size that does not match `n_fft`.
2158#[inline]
2159#[must_use]
2160pub fn make_window<T: Sample>(window: WindowType, n_fft: NonZeroUsize) -> NonEmptyVec<T> {
2161    let n_fft = n_fft.get();
2162    let mut w = vec![0.0; n_fft];
2163
2164    match window {
2165        WindowType::Rectangular => {
2166            w.fill(1.0);
2167        }
2168        WindowType::Hanning => {
2169            // Hann: 0.5 - 0.5*cos(2Ï€n/(N-1))
2170            let n1 = (n_fft - 1) as f64;
2171            for (n, v) in w.iter_mut().enumerate() {
2172                *v = 0.5f64.mul_add(-(2.0 * std::f64::consts::PI * (n as f64) / n1).cos(), 0.5);
2173            }
2174        }
2175        WindowType::Hamming => {
2176            // Hamming: 0.54 - 0.46*cos(2Ï€n/(N-1))
2177            let n1 = (n_fft - 1) as f64;
2178            for (n, v) in w.iter_mut().enumerate() {
2179                *v = 0.46f64.mul_add(-(2.0 * std::f64::consts::PI * (n as f64) / n1).cos(), 0.54);
2180            }
2181        }
2182        WindowType::Blackman => {
2183            // Blackman: 0.42 - 0.5*cos(2Ï€n/(N-1)) + 0.08*cos(4Ï€n/(N-1))
2184            let n1 = (n_fft - 1) as f64;
2185            for (n, v) in w.iter_mut().enumerate() {
2186                let a = 2.0 * std::f64::consts::PI * (n as f64) / n1;
2187                *v = 0.08f64.mul_add((2.0 * a).cos(), 0.5f64.mul_add(-a.cos(), 0.42));
2188            }
2189        }
2190        WindowType::Kaiser { beta } => {
2191            if n_fft == 1 {
2192                w[0] = 1.0;
2193            } else {
2194                let denom = modified_bessel_i0(beta);
2195                let n_max = (n_fft - 1) as f64 / 2.0;
2196
2197                for (i, value) in w.iter_mut().enumerate() {
2198                    let n = i as f64 - n_max;
2199                    let ratio = if n_max == 0.0 {
2200                        0.0
2201                    } else {
2202                        let normalized = n / n_max;
2203                        (1.0 - normalized * normalized).max(0.0)
2204                    };
2205                    let arg = beta * ratio.sqrt();
2206                    *value = if denom == 0.0 {
2207                        0.0
2208                    } else {
2209                        modified_bessel_i0(arg) / denom
2210                    };
2211                }
2212            }
2213        }
2214        WindowType::Gaussian { std } => (0..n_fft).for_each(|i| {
2215            let n = i as f64;
2216            let center: f64 = (n_fft - 1) as f64 / 2.0;
2217            let exponent: f64 = -0.5 * ((n - center) / std).powi(2);
2218            w[i] = exponent.exp();
2219        }),
2220        WindowType::Custom { coefficients, size } => {
2221            assert!(
2222                size.get() == n_fft,
2223                "Custom window size mismatch: expected {}, got {}. \
2224                 Custom windows must be pre-computed with the exact FFT size.",
2225                n_fft,
2226                size.get()
2227            );
2228            w.copy_from_slice(&coefficients);
2229        }
2230    }
2231
2232    // Coefficients are computed in f64 above, then mapped to the requested scalar.
2233    let w: Vec<T> = w.into_iter().map(T::from_f64).collect();
2234    // safety: window is guaranteed non-empty since n_fft > 0
2235    unsafe { NonEmptyVec::new_unchecked(w) }
2236}
2237
2238fn modified_bessel_i0(x: f64) -> f64 {
2239    let ax = x.abs();
2240    if ax <= 3.75 {
2241        let t = x / 3.75;
2242        let t2 = t * t;
2243        1.0 + t2
2244            * (3.515_622_9
2245                + t2 * (3.089_942_4
2246                    + t2 * (1.206_749_2
2247                        + t2 * (0.265_973_2 + t2 * (0.036_076_8 + t2 * 0.004_581_3)))))
2248    } else {
2249        let t = 3.75 / ax;
2250        let poly = 0.398_942_28
2251            + t * (0.013_285_92
2252                + t * (0.002_253_19
2253                    + t * (-0.001_575_65
2254                        + t * (0.009_162_81
2255                            + t * (-0.020_577_06
2256                                + t * (0.026_355_37 + t * (-0.016_476_33 + t * 0.003_923_77)))))));
2257
2258        (ax.exp() / (ax.sqrt() * (2.0 * std::f64::consts::PI).sqrt())) * poly
2259    }
2260}
2261
2262/// Convert Hz to mel scale using Slaney formula (librosa default, htk=False).
2263///
2264/// Uses a hybrid scale:
2265/// - Linear below 1000 Hz: mel = hz / (200/3)
2266/// - Logarithmic above 1000 Hz: mel = 15 + log(hz/1000) / log_step
2267///
2268/// This matches librosa's default behavior.
2269fn hz_to_mel(hz: f64) -> f64 {
2270    const F_MIN: f64 = 0.0;
2271    const F_SP: f64 = 200.0 / 3.0; // ~66.667
2272    const MIN_LOG_HZ: f64 = 1000.0;
2273    const MIN_LOG_MEL: f64 = (MIN_LOG_HZ - F_MIN) / F_SP; // = 15.0
2274    const LOGSTEP: f64 = 0.068_751_777_420_949_23; // ln(6.4) / 27
2275    if hz >= MIN_LOG_HZ {
2276        // Logarithmic region
2277        MIN_LOG_MEL + (hz / MIN_LOG_HZ).ln() / LOGSTEP
2278    } else {
2279        // Linear region
2280        (hz - F_MIN) / F_SP
2281    }
2282}
2283
2284/// Convert mel to Hz using Slaney formula (librosa default, htk=False).
2285///
2286/// Inverse of hz_to_mel.
2287fn mel_to_hz(mel: f64) -> f64 {
2288    const F_MIN: f64 = 0.0;
2289    const F_SP: f64 = 200.0 / 3.0; // ~66.667
2290    const MIN_LOG_HZ: f64 = 1000.0;
2291    const MIN_LOG_MEL: f64 = (MIN_LOG_HZ - F_MIN) / F_SP; // = 15.0
2292    const LOGSTEP: f64 = 0.068_751_777_420_949_23; // ln(6.4) / 27
2293
2294    if mel >= MIN_LOG_MEL {
2295        // Logarithmic region
2296        MIN_LOG_HZ * (LOGSTEP * (mel - MIN_LOG_MEL)).exp()
2297    } else {
2298        // Linear region
2299        F_SP.mul_add(mel, F_MIN)
2300    }
2301}
2302
2303fn build_mel_filterbank_matrix(
2304    sample_rate_hz: f64,
2305    n_fft: NonZeroUsize,
2306    n_mels: NonZeroUsize,
2307    f_min: f64,
2308    f_max: f64,
2309    norm: MelNorm,
2310) -> SpectrogramResult<SparseMatrix> {
2311    if sample_rate_hz <= 0.0 || !sample_rate_hz.is_finite() {
2312        return Err(SpectrogramError::invalid_input(
2313            "sample_rate_hz must be finite and > 0",
2314        ));
2315    }
2316    if f_min < 0.0 || f_min.is_infinite() {
2317        return Err(SpectrogramError::invalid_input("f_min must be >= 0"));
2318    }
2319    if f_max <= f_min {
2320        return Err(SpectrogramError::invalid_input("f_max must be > f_min"));
2321    }
2322    if f_max > sample_rate_hz * 0.5 {
2323        return Err(SpectrogramError::invalid_input("f_max must be <= Nyquist"));
2324    }
2325    let n_mels = n_mels.get();
2326    let n_fft = n_fft.get();
2327    let out_len = r2c_output_size(n_fft);
2328
2329    // FFT bin frequencies
2330    let df = sample_rate_hz / n_fft as f64;
2331
2332    // Mel points: n_mels + 2 (for triangular edges)
2333    let mel_min = hz_to_mel(f_min);
2334    let mel_max = hz_to_mel(f_max);
2335
2336    let n_points = n_mels + 2;
2337    let step = (mel_max - mel_min) / (n_points - 1) as f64;
2338
2339    let mut mel_points = Vec::with_capacity(n_points);
2340    for i in 0..n_points {
2341        mel_points.push((i as f64).mul_add(step, mel_min));
2342    }
2343
2344    let mut hz_points = Vec::with_capacity(n_points);
2345    for m in &mel_points {
2346        hz_points.push(mel_to_hz(*m));
2347    }
2348
2349    // Build filterbank as sparse matrix (librosa-style, in frequency space)
2350    // This builds triangular filters based on actual frequencies, not bin indices
2351    let mut fb = SparseMatrix::new(n_mels, out_len);
2352
2353    for m in 0..n_mels {
2354        let freq_left = hz_points[m];
2355        let freq_center = hz_points[m + 1];
2356        let freq_right = hz_points[m + 2];
2357
2358        let fdiff_left = freq_center - freq_left;
2359        let fdiff_right = freq_right - freq_center;
2360
2361        if fdiff_left == 0.0 || fdiff_right == 0.0 {
2362            // Degenerate triangle, skip
2363            continue;
2364        }
2365
2366        // For each FFT bin, compute the triangular weight based on its frequency
2367        for k in 0..out_len {
2368            let bin_freq = k as f64 * df;
2369
2370            // Lower slope: rises from freq_left to freq_center
2371            let lower = (bin_freq - freq_left) / fdiff_left;
2372
2373            // Upper slope: falls from freq_center to freq_right
2374            let upper = (freq_right - bin_freq) / fdiff_right;
2375
2376            // Triangle is the minimum of the two slopes, clipped to [0, 1]
2377            let weight = lower.min(upper).clamp(0.0, 1.0);
2378
2379            if weight > 0.0 {
2380                fb.set(m, k, weight);
2381            }
2382        }
2383    }
2384
2385    // Apply normalization
2386    match norm {
2387        MelNorm::None => {
2388            // No normalization needed
2389        }
2390        MelNorm::Slaney => {
2391            // Slaney-style area normalization: 2 / (hz_max - hz_min) for each triangle
2392            // NOTE: Uses Hz bandwidth, not mel bandwidth (to match librosa's implementation)
2393            for m in 0..n_mels {
2394                let mel_left = mel_points[m];
2395                let mel_right = mel_points[m + 2];
2396                let hz_left = mel_to_hz(mel_left);
2397                let hz_right = mel_to_hz(mel_right);
2398                let enorm = 2.0 / (hz_right - hz_left);
2399
2400                // Normalize all values in this row
2401                for val in &mut fb.values[m] {
2402                    *val *= enorm;
2403                }
2404            }
2405        }
2406        MelNorm::L1 => {
2407            // L1 normalization: sum of weights = 1.0
2408            for m in 0..n_mels {
2409                let sum: f64 = fb.values[m].iter().sum();
2410                if sum > 0.0 {
2411                    let normalizer = 1.0 / sum;
2412                    for val in &mut fb.values[m] {
2413                        *val *= normalizer;
2414                    }
2415                }
2416            }
2417        }
2418        MelNorm::L2 => {
2419            // L2 normalization: L2 norm = 1.0
2420            for m in 0..n_mels {
2421                let norm_val: f64 = fb.values[m].iter().map(|&v| v * v).sum::<f64>().sqrt();
2422                if norm_val > 0.0 {
2423                    let normalizer = 1.0 / norm_val;
2424                    for val in &mut fb.values[m] {
2425                        *val *= normalizer;
2426                    }
2427                }
2428            }
2429        }
2430    }
2431
2432    Ok(fb)
2433}
2434
2435/// Build a logarithmic frequency interpolation matrix.
2436///
2437/// Maps linearly-spaced FFT bins to logarithmically-spaced frequency bins
2438/// using linear interpolation.
2439fn build_loghz_matrix(
2440    sample_rate_hz: f64,
2441    n_fft: NonZeroUsize,
2442    n_bins: NonZeroUsize,
2443    f_min: f64,
2444    f_max: f64,
2445) -> SpectrogramResult<(SparseMatrix, NonEmptyVec<f64>)> {
2446    if sample_rate_hz <= 0.0 || !sample_rate_hz.is_finite() {
2447        return Err(SpectrogramError::invalid_input(
2448            "sample_rate_hz must be finite and > 0",
2449        ));
2450    }
2451    if f_min <= 0.0 || f_min.is_infinite() {
2452        return Err(SpectrogramError::invalid_input(
2453            "f_min must be finite and > 0",
2454        ));
2455    }
2456    if f_max <= f_min {
2457        return Err(SpectrogramError::invalid_input("f_max must be > f_min"));
2458    }
2459    if f_max > sample_rate_hz * 0.5 {
2460        return Err(SpectrogramError::invalid_input("f_max must be <= Nyquist"));
2461    }
2462
2463    let n_bins = n_bins.get();
2464    let n_fft = n_fft.get();
2465
2466    let out_len = r2c_output_size(n_fft);
2467    let df = sample_rate_hz / n_fft as f64;
2468
2469    // Generate logarithmically-spaced frequencies
2470    let log_f_min = f_min.ln();
2471    let log_f_max = f_max.ln();
2472    let log_step = (log_f_max - log_f_min) / (n_bins - 1) as f64;
2473
2474    let mut log_frequencies = Vec::with_capacity(n_bins);
2475    for i in 0..n_bins {
2476        let log_f = (i as f64).mul_add(log_step, log_f_min);
2477        log_frequencies.push(log_f.exp());
2478    }
2479    // safety: n_bins > 0
2480    let log_frequencies = unsafe { NonEmptyVec::new_unchecked(log_frequencies) };
2481
2482    // Build interpolation matrix as sparse matrix
2483    let mut matrix = SparseMatrix::new(n_bins, out_len);
2484
2485    for (bin_idx, &target_freq) in log_frequencies.iter().enumerate() {
2486        // Find the two FFT bins that bracket this frequency
2487        let exact_bin = target_freq / df;
2488        let lower_bin = exact_bin.floor() as usize;
2489        let upper_bin = (exact_bin.ceil() as usize).min(out_len - 1);
2490
2491        if lower_bin >= out_len {
2492            continue;
2493        }
2494
2495        if lower_bin == upper_bin {
2496            // Exact match
2497            matrix.set(bin_idx, lower_bin, 1.0);
2498        } else {
2499            // Linear interpolation
2500            let frac = exact_bin - lower_bin as f64;
2501            matrix.set(bin_idx, lower_bin, 1.0 - frac);
2502            if upper_bin < out_len {
2503                matrix.set(bin_idx, upper_bin, frac);
2504            }
2505        }
2506    }
2507
2508    Ok((matrix, log_frequencies))
2509}
2510
2511fn mel_band_centres_hz(
2512    n_mels: NonZeroUsize,
2513    sample_rate_hz: f64,
2514    nyquist_hz: f64,
2515) -> NonEmptyVec<f64> {
2516    let f_min = 0.0;
2517    let f_max = nyquist_hz.min(sample_rate_hz * 0.5);
2518
2519    let mel_min = hz_to_mel(f_min);
2520    let mel_max = hz_to_mel(f_max);
2521    let n_mels = n_mels.get();
2522    let step = (mel_max - mel_min) / (n_mels + 1) as f64;
2523
2524    let mut centres = Vec::with_capacity(n_mels);
2525    for i in 0..n_mels {
2526        let mel = (i as f64 + 1.0).mul_add(step, mel_min);
2527        centres.push(mel_to_hz(mel));
2528    }
2529    // safety: centres is guaranteed non-empty since n_mels > 0
2530    unsafe { NonEmptyVec::new_unchecked(centres) }
2531}
2532
2533/// Spectrogram structure holding the computed spectrogram data and metadata.
2534///
2535/// # Type Parameters
2536///
2537/// * `FreqScale`: The frequency scale type (e.g., `LinearHz`, `Mel`, `LogHz`, etc.).
2538/// * `AmpScale`: The amplitude scale type (e.g., `Power`, `Magnitude`, `Decibels`).
2539///
2540/// # Fields
2541///
2542/// * `data`: A 2D array containing the spectrogram data.
2543/// * `axes`: The axes of the spectrogram (frequency and time).
2544/// * `params`: The parameters used to compute the spectrogram.
2545/// * `_amp`: A phantom data marker for the amplitude scale type.
2546#[derive(Debug, Clone)]
2547#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
2548pub struct Spectrogram<FreqScale, AmpScale, T = f64>
2549where
2550    AmpScale: AmpScaleSpec + 'static,
2551    FreqScale: Copy + Clone + 'static,
2552{
2553    data: Array2<T>,
2554    axes: Axes<FreqScale>,
2555    params: SpectrogramParams,
2556    #[cfg_attr(feature = "serde", serde(skip))]
2557    _amp: PhantomData<AmpScale>,
2558}
2559
2560impl<FreqScale, AmpScale, T> Spectrogram<FreqScale, AmpScale, T>
2561where
2562    AmpScale: AmpScaleSpec + 'static,
2563    FreqScale: Copy + Clone + 'static,
2564    T: Sample,
2565{
2566    /// Get the X-axis label.
2567    ///
2568    /// # Returns
2569    ///
2570    /// A static string slice representing the X-axis label.
2571    #[inline]
2572    #[must_use]
2573    pub const fn x_axis_label() -> &'static str {
2574        "Time (s)"
2575    }
2576
2577    /// Get the Y-axis label based on the frequency scale type.
2578    ///
2579    /// # Returns
2580    ///
2581    /// A static string slice representing the Y-axis label.
2582    #[inline]
2583    #[must_use]
2584    pub fn y_axis_label() -> &'static str {
2585        match std::any::TypeId::of::<FreqScale>() {
2586            id if id == std::any::TypeId::of::<LinearHz>() => "Frequency (Hz)",
2587            id if id == std::any::TypeId::of::<Mel>() => "Frequency (Mel)",
2588            id if id == std::any::TypeId::of::<LogHz>() => "Frequency (Log Hz)",
2589            id if id == std::any::TypeId::of::<Erb>() => "Frequency (ERB)",
2590            id if id == std::any::TypeId::of::<Cqt>() => "Frequency (CQT Bins)",
2591            _ => "Frequency",
2592        }
2593    }
2594
2595    /// Internal constructor. Only callable inside the crate.
2596    ///
2597    /// All inputs must already be validated and consistent.
2598    pub(crate) fn new(data: Array2<T>, axes: Axes<FreqScale>, params: SpectrogramParams) -> Self {
2599        debug_assert_eq!(data.nrows(), axes.frequencies().len().get());
2600        debug_assert_eq!(data.ncols(), axes.times().len().get());
2601
2602        Self {
2603            data,
2604            axes,
2605            params,
2606            _amp: PhantomData,
2607        }
2608    }
2609
2610    /// Set spectrogram data matrix
2611    ///
2612    /// # Arguments
2613    ///
2614    /// * `data` - The new spectrogram data matrix.
2615    #[inline]
2616    pub fn set_data(&mut self, data: Array2<T>) {
2617        self.data = data;
2618    }
2619
2620    /// Spectrogram data matrix
2621    ///
2622    /// # Returns
2623    ///
2624    /// A reference to the spectrogram data matrix.
2625    #[inline]
2626    #[must_use]
2627    pub const fn data(&self) -> &Array2<T> {
2628        &self.data
2629    }
2630
2631    /// Consume the spectrogram and extract the data matrix.
2632    ///
2633    /// This method moves the data out of the spectrogram, consuming it.
2634    /// Useful for transferring ownership to Python without copying.
2635    ///
2636    /// # Returns
2637    ///
2638    /// The owned spectrogram data matrix.
2639    #[inline]
2640    #[must_use]
2641    pub fn into_data(self) -> Array2<T> {
2642        self.data
2643    }
2644
2645    /// Axes of the spectrogram
2646    ///
2647    /// # Returns
2648    ///
2649    /// A reference to the axes of the spectrogram.
2650    #[inline]
2651    #[must_use]
2652    pub const fn axes(&self) -> &Axes<FreqScale> {
2653        &self.axes
2654    }
2655
2656    /// Frequency axis in Hz
2657    ///
2658    /// # Returns
2659    ///
2660    /// A reference to the frequency axis in Hz.
2661    #[inline]
2662    #[must_use]
2663    pub fn frequencies(&self) -> &NonEmptySlice<f64> {
2664        self.axes.frequencies()
2665    }
2666
2667    /// Frequency range in Hz (min, max)
2668    ///
2669    /// # Returns
2670    ///
2671    /// A tuple containing the minimum and maximum frequencies in Hz.
2672    #[inline]
2673    #[must_use]
2674    pub const fn frequency_range(&self) -> (f64, f64) {
2675        self.axes.frequency_range()
2676    }
2677
2678    /// Time axis in seconds
2679    ///
2680    /// # Returns
2681    ///
2682    /// A reference to the time axis in seconds.
2683    #[inline]
2684    #[must_use]
2685    pub fn times(&self) -> &NonEmptySlice<f64> {
2686        self.axes.times()
2687    }
2688
2689    /// Spectrogram computation parameters
2690    ///
2691    /// # Returns
2692    ///
2693    /// A reference to the spectrogram computation parameters.
2694    #[inline]
2695    #[must_use]
2696    pub const fn params(&self) -> &SpectrogramParams {
2697        &self.params
2698    }
2699
2700    /// Duration of the spectrogram in seconds
2701    ///
2702    /// # Returns
2703    ///
2704    /// The duration of the spectrogram in seconds.
2705    #[inline]
2706    #[must_use]
2707    pub fn duration(&self) -> f64 {
2708        self.axes.duration()
2709    }
2710
2711    /// If this is a dB spectrogram, return the (min, max) dB values. otherwise do the maths to compute dB range.
2712    ///
2713    /// # Returns
2714    ///
2715    /// The (min, max) dB values of the spectrogram, or `None` if the amplitude scale is unknown.
2716    #[inline]
2717    #[must_use]
2718    pub fn db_range(&self) -> Option<(f64, f64)> {
2719        let type_self = std::any::TypeId::of::<AmpScale>();
2720
2721        if type_self == std::any::TypeId::of::<Decibels>() {
2722            let mut min = f64::INFINITY;
2723            let mut max = f64::NEG_INFINITY;
2724            for &v in &self.data {
2725                let v = v.to_f64().unwrap_or(0.0);
2726                if v < min {
2727                    min = v;
2728                }
2729                if v > max {
2730                    max = v;
2731                }
2732            }
2733            Some((min, max))
2734        } else if type_self == std::any::TypeId::of::<Power>() {
2735            // Not a dB spectrogram; compute dB range from power values
2736            let mut min_db = f64::INFINITY;
2737            let mut max_db = f64::NEG_INFINITY;
2738            for &v in &self.data {
2739                let db = 10.0 * (v.to_f64().unwrap_or(0.0) + EPS).log10();
2740                if db < min_db {
2741                    min_db = db;
2742                }
2743                if db > max_db {
2744                    max_db = db;
2745                }
2746            }
2747            Some((min_db, max_db))
2748        } else if type_self == std::any::TypeId::of::<Magnitude>() {
2749            // Not a dB spectrogram; compute dB range from magnitude values
2750            let mut min_db = f64::INFINITY;
2751            let mut max_db = f64::NEG_INFINITY;
2752
2753            for &v in &self.data {
2754                let v = v.to_f64().unwrap_or(0.0);
2755                let power = v * v;
2756                let db = 10.0 * (power + EPS).log10();
2757                if db < min_db {
2758                    min_db = db;
2759                }
2760                if db > max_db {
2761                    max_db = db;
2762                }
2763            }
2764
2765            Some((min_db, max_db))
2766        } else {
2767            // Unknown AmpScale type; return dummy values
2768            None
2769        }
2770    }
2771
2772    /// Number of frequency bins
2773    ///
2774    /// # Returns
2775    ///
2776    /// The number of frequency bins in the spectrogram.
2777    #[inline]
2778    #[must_use]
2779    pub fn n_bins(&self) -> NonZeroUsize {
2780        // safety: data.nrows() > 0 is guaranteed by construction
2781        unsafe { NonZeroUsize::new_unchecked(self.data.nrows()) }
2782    }
2783
2784    /// Number of time frames in the spectrogram
2785    ///
2786    /// # Returns
2787    ///
2788    /// The number of time frames (columns) in the spectrogram.
2789    #[inline]
2790    #[must_use]
2791    pub fn n_frames(&self) -> NonZeroUsize {
2792        // safety: data.ncols() > 0 is guaranteed by construction
2793        unsafe { NonZeroUsize::new_unchecked(self.data.ncols()) }
2794    }
2795}
2796
2797impl<FreqScale, AmpScale, T> AsRef<Array2<T>> for Spectrogram<FreqScale, AmpScale, T>
2798where
2799    FreqScale: Copy + Clone + 'static,
2800    AmpScale: AmpScaleSpec + 'static,
2801    T: Sample,
2802{
2803    #[inline]
2804    fn as_ref(&self) -> &Array2<T> {
2805        &self.data
2806    }
2807}
2808
2809impl<FreqScale, AmpScale, T> Deref for Spectrogram<FreqScale, AmpScale, T>
2810where
2811    FreqScale: Copy + Clone + 'static,
2812    AmpScale: AmpScaleSpec + 'static,
2813    T: Sample,
2814{
2815    type Target = Array2<T>;
2816
2817    #[inline]
2818    fn deref(&self) -> &Self::Target {
2819        &self.data
2820    }
2821}
2822
2823impl<FreqScale, AmpScale, T> DerefMut for Spectrogram<FreqScale, AmpScale, T>
2824where
2825    FreqScale: Copy + Clone + 'static,
2826    AmpScale: AmpScaleSpec + 'static,
2827    T: Sample,
2828{
2829    #[inline]
2830    fn deref_mut(&mut self) -> &mut Self::Target {
2831        &mut self.data
2832    }
2833}
2834
2835impl<AmpScale, T: Sample> Spectrogram<LinearHz, AmpScale, T>
2836where
2837    AmpScale: AmpScaleSpec + 'static,
2838{
2839    /// Compute a linear-frequency spectrogram from audio samples.
2840    ///
2841    /// This is a convenience method that creates a planner internally and computes
2842    /// the spectrogram in one call. For processing multiple signals with the same
2843    /// parameters, use [`SpectrogramPlanner::linear_plan`] to create a reusable plan.
2844    ///
2845    /// # Arguments
2846    ///
2847    /// * `samples` - Audio samples (any type that can be converted to a slice)
2848    /// * `params` - Spectrogram computation parameters
2849    /// * `db` - Optional logarithmic scaling parameters (only used when `AmpScale = Decibels`)
2850    ///
2851    /// # Returns
2852    ///
2853    /// A linear-frequency spectrogram with the specified amplitude scale.
2854    ///
2855    /// # Errors
2856    ///
2857    /// Returns an error if:
2858    /// - The samples slice is empty
2859    /// - Parameters are invalid
2860    /// - FFT computation fails
2861    ///
2862    /// # Examples
2863    ///
2864    /// ```
2865    /// use spectrograms::*;
2866    /// use non_empty_slice::non_empty_vec;
2867    ///
2868    /// # fn example() -> SpectrogramResult<()> {
2869    /// // Create a simple test signal
2870    /// let sample_rate = 16000.0;
2871    /// let samples_vec: Vec<f64> = (0..16000).map(|i| {
2872    ///     (2.0 * std::f64::consts::PI * 440.0 * i as f64 / sample_rate).sin()
2873    /// }).collect();
2874    /// let samples = non_empty_slice::NonEmptyVec::new(samples_vec).unwrap();
2875    ///
2876    /// // Set up parameters
2877    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
2878    /// let params = SpectrogramParams::new(stft, sample_rate)?;
2879    ///
2880    /// // Compute power spectrogram
2881    /// let spec = LinearPowerSpectrogram::compute(&samples, &params, None)?;
2882    ///
2883    /// println!("Computed spectrogram: {} bins x {} frames", spec.n_bins(), spec.n_frames());
2884    /// # Ok(())
2885    /// # }
2886    /// ```
2887    #[inline]
2888    pub fn compute(
2889        samples: &NonEmptySlice<T>,
2890        params: &SpectrogramParams,
2891        db: Option<&LogParams>,
2892    ) -> SpectrogramResult<Self> {
2893        let planner = SpectrogramPlanner::new();
2894        let mut plan = planner.linear_plan(params, db)?;
2895        plan.compute(samples)
2896    }
2897}
2898
2899impl<AmpScale, T: Sample> Spectrogram<Mel, AmpScale, T>
2900where
2901    AmpScale: AmpScaleSpec + 'static,
2902{
2903    /// Compute a mel-frequency spectrogram from audio samples.
2904    ///
2905    /// This is a convenience method that creates a planner internally and computes
2906    /// the spectrogram in one call. For processing multiple signals with the same
2907    /// parameters, use [`SpectrogramPlanner::mel_plan`] to create a reusable plan.
2908    ///
2909    /// # Arguments
2910    ///
2911    /// * `samples` - Audio samples (any type that can be converted to a slice)
2912    /// * `params` - Spectrogram computation parameters
2913    /// * `mel` - Mel filterbank parameters
2914    /// * `db` - Optional logarithmic scaling parameters (only used when `AmpScale = Decibels`)
2915    ///
2916    /// # Returns
2917    ///
2918    /// A mel-frequency spectrogram with the specified amplitude scale.
2919    ///
2920    /// # Errors
2921    ///
2922    /// Returns an error if:
2923    /// - The samples slice is empty
2924    /// - Parameters are invalid
2925    /// - Mel `f_max` exceeds Nyquist frequency
2926    /// - FFT computation fails
2927    ///
2928    /// # Examples
2929    ///
2930    /// ```
2931    /// use spectrograms::*;
2932    /// use non_empty_slice::non_empty_vec;
2933    ///
2934    /// # fn example() -> SpectrogramResult<()> {
2935    /// // Create a simple test signal
2936    /// let sample_rate = 16000.0;
2937    /// let samples_vec: Vec<f64> = (0..16000).map(|i| {
2938    ///     (2.0 * std::f64::consts::PI * 440.0 * i as f64 / sample_rate).sin()
2939    /// }).collect();
2940    /// let samples = non_empty_slice::NonEmptyVec::new(samples_vec).unwrap();
2941    ///
2942    /// // Set up parameters
2943    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
2944    /// let params = SpectrogramParams::new(stft, sample_rate)?;
2945    /// let mel = MelParams::new(nzu!(80), 0.0, 8000.0)?;
2946    ///
2947    /// // Compute mel spectrogram in dB scale
2948    /// let db = LogParams::new(-80.0)?;
2949    /// let spec = MelDbSpectrogram::compute(&samples, &params, &mel, Some(&db))?;
2950    ///
2951    /// println!("Computed mel spectrogram: {} mels x {} frames", spec.n_bins(), spec.n_frames());
2952    /// # Ok(())
2953    /// # }
2954    /// ```
2955    #[inline]
2956    pub fn compute(
2957        samples: &NonEmptySlice<T>,
2958        params: &SpectrogramParams,
2959        mel: &MelParams,
2960        db: Option<&LogParams>,
2961    ) -> SpectrogramResult<Self> {
2962        let planner = SpectrogramPlanner::new();
2963        let mut plan = planner.mel_plan(params, mel, db)?;
2964        plan.compute(samples)
2965    }
2966}
2967
2968impl<AmpScale, T: Sample> Spectrogram<Erb, AmpScale, T>
2969where
2970    AmpScale: AmpScaleSpec + 'static,
2971{
2972    /// Compute an ERB-frequency spectrogram from audio samples.
2973    ///
2974    /// This is a convenience method that creates a planner internally and computes
2975    /// the spectrogram in one call. For processing multiple signals with the same
2976    /// parameters, use [`SpectrogramPlanner::erb_plan`] to create a reusable plan.
2977    ///
2978    /// # Arguments
2979    ///
2980    /// * `samples` - Audio samples (any type that can be converted to a slice)
2981    /// * `params` - Spectrogram computation parameters
2982    /// * `erb` - ERB frequency scale parameters
2983    /// * `db` - Optional logarithmic scaling parameters (only used when `AmpScale = Decibels`)
2984    ///
2985    /// # Returns
2986    ///
2987    /// An ERB-scale spectrogram with the specified amplitude scale.
2988    ///
2989    /// # Errors
2990    ///
2991    /// Returns an error if:
2992    /// - The samples slice is empty
2993    /// - Parameters are invalid
2994    /// - FFT computation fails
2995    ///
2996    /// # Examples
2997    ///
2998    /// ```
2999    /// use spectrograms::*;
3000    /// use non_empty_slice::non_empty_vec;
3001    ///
3002    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
3003    /// let samples = non_empty_vec![0.0; nzu!(16000)];
3004    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
3005    /// let params = SpectrogramParams::new(stft, 16000.0)?;
3006    /// let erb = ErbParams::speech_standard();
3007    ///
3008    /// let spec = ErbPowerSpectrogram::compute(&samples, &params, &erb, None)?;
3009    /// assert_eq!(spec.n_bins(), nzu!(40));
3010    /// # Ok(())
3011    /// # }
3012    /// ```
3013    #[inline]
3014    pub fn compute(
3015        samples: &NonEmptySlice<T>,
3016        params: &SpectrogramParams,
3017        erb: &ErbParams,
3018        db: Option<&LogParams>,
3019    ) -> SpectrogramResult<Self> {
3020        let planner = SpectrogramPlanner::new();
3021        let mut plan = planner.erb_plan(params, erb, db)?;
3022        plan.compute(samples)
3023    }
3024}
3025
3026impl<AmpScale, T: Sample> Spectrogram<LogHz, AmpScale, T>
3027where
3028    AmpScale: AmpScaleSpec + 'static,
3029{
3030    /// Compute a logarithmic-frequency spectrogram from audio samples.
3031    ///
3032    /// This is a convenience method that creates a planner internally and computes
3033    /// the spectrogram in one call. For processing multiple signals with the same
3034    /// parameters, use [`SpectrogramPlanner::log_hz_plan`] to create a reusable plan.
3035    ///
3036    /// # Arguments
3037    ///
3038    /// * `samples` - Audio samples (any type that can be converted to a slice)
3039    /// * `params` - Spectrogram computation parameters
3040    /// * `loghz` - Logarithmic frequency scale parameters
3041    /// * `db` - Optional logarithmic scaling parameters (only used when `AmpScale = Decibels`)
3042    ///
3043    /// # Returns
3044    ///
3045    /// A logarithmic-frequency spectrogram with the specified amplitude scale.
3046    ///
3047    /// # Errors
3048    ///
3049    /// Returns an error if:
3050    /// - The samples slice is empty
3051    /// - Parameters are invalid
3052    /// - FFT computation fails
3053    ///
3054    /// # Examples
3055    ///
3056    /// ```
3057    /// use spectrograms::*;
3058    /// use non_empty_slice::non_empty_vec;
3059    ///
3060    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
3061    /// let samples = non_empty_vec![0.0; nzu!(16000)];
3062    /// let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true)?;
3063    /// let params = SpectrogramParams::new(stft, 16000.0)?;
3064    /// let loghz = LogHzParams::new(nzu!(128), 20.0, 8000.0)?;
3065    ///
3066    /// let spec = LogHzPowerSpectrogram::compute(&samples, &params, &loghz, None)?;
3067    /// assert_eq!(spec.n_bins(), nzu!(128));
3068    /// # Ok(())
3069    /// # }
3070    /// ```
3071    #[inline]
3072    pub fn compute(
3073        samples: &NonEmptySlice<T>,
3074        params: &SpectrogramParams,
3075        loghz: &LogHzParams,
3076        db: Option<&LogParams>,
3077    ) -> SpectrogramResult<Self> {
3078        let planner = SpectrogramPlanner::new();
3079        let mut plan = planner.log_hz_plan(params, loghz, db)?;
3080        plan.compute(samples)
3081    }
3082}
3083
3084impl<AmpScale, T: Sample> Spectrogram<Cqt, AmpScale, T>
3085where
3086    AmpScale: AmpScaleSpec + 'static,
3087{
3088    /// Compute a constant-Q transform (CQT) spectrogram from audio samples.
3089    ///
3090    /// # Arguments
3091    ///
3092    /// * `samples` - Audio samples (any type that can be converted to a slice)
3093    /// * `params` - Spectrogram computation parameters
3094    /// * `cqt` - CQT parameters
3095    /// * `db` - Optional logarithmic scaling parameters (only used when `AmpScale = Decibels`)
3096    ///
3097    /// # Returns
3098    ///
3099    /// A CQT spectrogram with the specified amplitude scale.
3100    ///
3101    /// # Errors
3102    ///
3103    /// Returns an error if:
3104    ///
3105    #[inline]
3106    pub fn compute(
3107        samples: &NonEmptySlice<T>,
3108        params: &SpectrogramParams,
3109        cqt: &CqtParams,
3110        db: Option<&LogParams>,
3111    ) -> SpectrogramResult<Self> {
3112        let planner = SpectrogramPlanner::new();
3113        let mut plan = planner.cqt_plan(params, cqt, db)?;
3114        plan.compute(samples)
3115    }
3116}
3117
3118// ========================
3119// Display implementations
3120// ========================
3121
3122/// Helper function to get amplitude scale name
3123fn amp_scale_name<AmpScale>() -> &'static str
3124where
3125    AmpScale: AmpScaleSpec + 'static,
3126{
3127    match std::any::TypeId::of::<AmpScale>() {
3128        id if id == std::any::TypeId::of::<Power>() => "Power",
3129        id if id == std::any::TypeId::of::<Magnitude>() => "Magnitude",
3130        id if id == std::any::TypeId::of::<Decibels>() => "Decibels",
3131        _ => "Unknown",
3132    }
3133}
3134
3135/// Helper function to get frequency scale name
3136fn freq_scale_name<FreqScale>() -> &'static str
3137where
3138    FreqScale: Copy + Clone + 'static,
3139{
3140    match std::any::TypeId::of::<FreqScale>() {
3141        id if id == std::any::TypeId::of::<LinearHz>() => "Linear Hz",
3142        id if id == std::any::TypeId::of::<LogHz>() => "Log Hz",
3143        id if id == std::any::TypeId::of::<Mel>() => "Mel",
3144        id if id == std::any::TypeId::of::<Erb>() => "ERB",
3145        id if id == std::any::TypeId::of::<Cqt>() => "CQT",
3146        _ => "Unknown",
3147    }
3148}
3149
3150impl<FreqScale, AmpScale, T> core::fmt::Display for Spectrogram<FreqScale, AmpScale, T>
3151where
3152    AmpScale: AmpScaleSpec + 'static,
3153    FreqScale: Copy + Clone + 'static,
3154    T: Sample,
3155{
3156    #[inline]
3157    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
3158        let (freq_min, freq_max) = self.frequency_range();
3159        let duration = self.duration();
3160        let (rows, cols) = self.data.dim();
3161
3162        // Alternative formatting (#) provides more detailed output with data
3163        if f.alternate() {
3164            writeln!(f, "Spectrogram {{")?;
3165            writeln!(f, "  Frequency Scale: {}", freq_scale_name::<FreqScale>())?;
3166            writeln!(f, "  Amplitude Scale: {}", amp_scale_name::<AmpScale>())?;
3167            writeln!(f, "  Shape: {rows} frequency bins × {cols} time frames")?;
3168            writeln!(f, "  Frequency Range: {freq_min:.2} Hz - {freq_max:.2} Hz")?;
3169            writeln!(f, "  Duration: {duration:.3} s")?;
3170            writeln!(f)?;
3171
3172            // Display parameters
3173            writeln!(f, "  Parameters:")?;
3174            writeln!(f, "    Sample Rate: {} Hz", self.params.sample_rate_hz())?;
3175            writeln!(f, "    FFT Size: {}", self.params.stft().n_fft())?;
3176            writeln!(f, "    Hop Size: {}", self.params.stft().hop_size())?;
3177            writeln!(f, "    Window: {:?}", self.params.stft().window())?;
3178            writeln!(f, "    Centered: {}", self.params.stft().centre())?;
3179            writeln!(f)?;
3180
3181            // Display data statistics (converted to f64 for formatting)
3182            let data_slice: Vec<f64> = self
3183                .data
3184                .as_slice()
3185                .unwrap_or(&[])
3186                .iter()
3187                .map(|&v| v.to_f64().unwrap_or(0.0))
3188                .collect();
3189            if !data_slice.is_empty() {
3190                let (min_val, max_val) = min_max_single_pass(&data_slice);
3191                let mean = data_slice.iter().sum::<f64>() / data_slice.len() as f64;
3192                writeln!(f, "  Data Statistics:")?;
3193                writeln!(f, "    Min: {min_val:.6}")?;
3194                writeln!(f, "    Max: {max_val:.6}")?;
3195                writeln!(f, "    Mean: {mean:.6}")?;
3196                writeln!(f)?;
3197            }
3198
3199            // Display actual data (truncated if too large)
3200            writeln!(f, "  Data Matrix:")?;
3201            let max_rows_to_display = 5;
3202            let max_cols_to_display = 5;
3203
3204            for i in 0..rows.min(max_rows_to_display) {
3205                write!(f, "    [")?;
3206                for j in 0..cols.min(max_cols_to_display) {
3207                    if j > 0 {
3208                        write!(f, ", ")?;
3209                    }
3210                    write!(f, "{:9.4}", self.data[[i, j]].to_f64().unwrap_or(0.0))?;
3211                }
3212                if cols > max_cols_to_display {
3213                    write!(f, ", ... ({} more)", cols - max_cols_to_display)?;
3214                }
3215                writeln!(f, "]")?;
3216            }
3217
3218            if rows > max_rows_to_display {
3219                writeln!(f, "    ... ({} more rows)", rows - max_rows_to_display)?;
3220            }
3221
3222            write!(f, "}}")?;
3223        } else {
3224            // Default formatting: compact summary
3225            write!(
3226                f,
3227                "Spectrogram<{}, {}>[{}x{}] ({:.2}-{:.2} Hz, {:.3}s)",
3228                freq_scale_name::<FreqScale>(),
3229                amp_scale_name::<AmpScale>(),
3230                rows,
3231                cols,
3232                freq_min,
3233                freq_max,
3234                duration
3235            )?;
3236        }
3237
3238        Ok(())
3239    }
3240}
3241
3242#[derive(Debug, Clone)]
3243#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3244pub struct FrequencyAxis<FreqScale>
3245where
3246    FreqScale: Copy + Clone + 'static,
3247{
3248    frequencies: NonEmptyVec<f64>,
3249    #[cfg_attr(feature = "serde", serde(skip))]
3250    _marker: PhantomData<FreqScale>,
3251}
3252
3253impl<FreqScale> FrequencyAxis<FreqScale>
3254where
3255    FreqScale: Copy + Clone + 'static,
3256{
3257    pub(crate) const fn new(frequencies: NonEmptyVec<f64>) -> Self {
3258        Self {
3259            frequencies,
3260            _marker: PhantomData,
3261        }
3262    }
3263
3264    /// Get the frequency values in Hz.
3265    ///
3266    /// # Returns
3267    ///
3268    /// Returns a non-empty slice of frequencies.
3269    #[inline]
3270    #[must_use]
3271    pub fn frequencies(&self) -> &NonEmptySlice<f64> {
3272        &self.frequencies
3273    }
3274
3275    /// Get the frequency range (min, max) in Hz.
3276    ///
3277    /// # Returns
3278    ///
3279    /// Returns a tuple containing the minimum and maximum frequency.
3280    #[inline]
3281    #[must_use]
3282    pub const fn frequency_range(&self) -> (f64, f64) {
3283        let data = self.frequencies.as_slice();
3284        let min = data[0];
3285        let max_idx = data.len().saturating_sub(1); // safe for non-empty
3286        let max = data[max_idx];
3287        (min, max)
3288    }
3289
3290    /// Get the number of frequency bins.
3291    ///
3292    /// # Returns
3293    ///
3294    /// Returns the number of frequency bins as a NonZeroUsize.
3295    #[inline]
3296    #[must_use]
3297    pub const fn len(&self) -> NonZeroUsize {
3298        self.frequencies.len()
3299    }
3300}
3301
3302/// Spectrogram axes container.
3303///
3304/// Holds frequency and time axes.
3305#[derive(Debug, Clone)]
3306#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3307pub struct Axes<FreqScale>
3308where
3309    FreqScale: Copy + Clone + 'static,
3310{
3311    freq: FrequencyAxis<FreqScale>,
3312    times: NonEmptyVec<f64>,
3313}
3314
3315impl<FreqScale> Axes<FreqScale>
3316where
3317    FreqScale: Copy + Clone + 'static,
3318{
3319    pub(crate) const fn new(freq: FrequencyAxis<FreqScale>, times: NonEmptyVec<f64>) -> Self {
3320        Self { freq, times }
3321    }
3322
3323    /// Get the frequency values in Hz.
3324    ///
3325    /// # Returns
3326    ///
3327    /// Returns a non-empty slice of frequencies.
3328    #[inline]
3329    #[must_use]
3330    pub fn frequencies(&self) -> &NonEmptySlice<f64> {
3331        self.freq.frequencies()
3332    }
3333
3334    /// Get the time values in seconds.
3335    ///
3336    /// # Returns
3337    ///
3338    /// Returns a non-empty slice of time values.
3339    #[inline]
3340    #[must_use]
3341    pub fn times(&self) -> &NonEmptySlice<f64> {
3342        &self.times
3343    }
3344
3345    /// Get the frequency range (min, max) in Hz.
3346    ///
3347    /// # Returns
3348    ///
3349    /// Returns a tuple containing the minimum and maximum frequency.
3350    #[inline]
3351    #[must_use]
3352    pub const fn frequency_range(&self) -> (f64, f64) {
3353        self.freq.frequency_range()
3354    }
3355
3356    /// Get the duration of the spectrogram in seconds.
3357    ///
3358    /// # Returns
3359    ///
3360    /// Returns the duration in seconds.
3361    #[inline]
3362    #[must_use]
3363    pub fn duration(&self) -> f64 {
3364        *self.times.last()
3365    }
3366}
3367
3368// Enum types for frequency and amplitude scales
3369
3370/// Linear frequency scale
3371#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3372#[cfg_attr(feature = "python", pyclass(from_py_object))]
3373#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3374#[non_exhaustive]
3375pub enum LinearHz {
3376    _Phantom,
3377}
3378
3379/// Logarithmic frequency scale
3380#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3381#[cfg_attr(feature = "python", pyclass(from_py_object))]
3382#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3383#[non_exhaustive]
3384pub enum LogHz {
3385    _Phantom,
3386}
3387
3388/// Mel frequency scale
3389#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3390#[cfg_attr(feature = "python", pyclass(from_py_object))]
3391#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3392#[non_exhaustive]
3393pub enum Mel {
3394    _Phantom,
3395}
3396
3397/// ERB/gammatone frequency scale
3398#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3399#[cfg_attr(feature = "python", pyclass(from_py_object))]
3400#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3401#[non_exhaustive]
3402pub enum Erb {
3403    _Phantom,
3404}
3405pub type Gammatone = Erb;
3406
3407/// Constant-Q Transform frequency scale
3408#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3409#[cfg_attr(feature = "python", pyclass(from_py_object))]
3410#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3411#[non_exhaustive]
3412pub enum Cqt {
3413    _Phantom,
3414}
3415
3416// Amplitude scales
3417
3418/// Power amplitude scale
3419#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3420#[cfg_attr(feature = "python", pyclass(from_py_object))]
3421#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3422#[non_exhaustive]
3423pub enum Power {
3424    _Phantom,
3425}
3426
3427/// Decibel amplitude scale
3428#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3429#[cfg_attr(feature = "python", pyclass(from_py_object))]
3430#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3431#[non_exhaustive]
3432pub enum Decibels {
3433    _Phantom,
3434}
3435
3436/// Magnitude amplitude scale
3437#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
3438#[cfg_attr(feature = "python", pyclass(from_py_object))]
3439#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3440#[non_exhaustive]
3441pub enum Magnitude {
3442    _Phantom,
3443}
3444
3445/// STFT parameters for spectrogram computation.
3446///
3447/// * `n_fft`: Size of the FFT window.
3448/// * `hop_size`: Number of samples between successive frames.
3449/// * window: Window function to apply to each frame.
3450/// * centre: Whether to pad the input signal so that frames are centered.
3451#[derive(Debug, Clone, PartialEq)]
3452#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3453pub struct StftParams {
3454    n_fft: NonZeroUsize,
3455    hop_size: NonZeroUsize,
3456    window: WindowType,
3457    centre: bool,
3458}
3459
3460impl StftParams {
3461    /// Create new STFT parameters.
3462    ///
3463    /// # Arguments
3464    ///
3465    /// * `n_fft` - Size of the FFT window
3466    /// * `hop_size` - Number of samples between successive frames
3467    /// * `window` - Window function to apply to each frame
3468    /// * `centre` - Whether to pad the input signal so that frames are centered
3469    ///
3470    /// # Errors
3471    ///
3472    /// Returns an error if:
3473    /// - `hop_size` > `n_fft`
3474    /// - Custom window size doesn't match `n_fft`
3475    ///
3476    /// # Returns
3477    ///
3478    /// New `StftParams` instance.
3479    #[inline]
3480    pub fn new(
3481        n_fft: NonZeroUsize,
3482        hop_size: NonZeroUsize,
3483        window: WindowType,
3484        centre: bool,
3485    ) -> SpectrogramResult<Self> {
3486        if hop_size.get() > n_fft.get() {
3487            return Err(SpectrogramError::invalid_input("hop_size must be <= n_fft"));
3488        }
3489
3490        // Validate custom window size matches n_fft
3491        if let WindowType::Custom { size, .. } = &window {
3492            if size.get() != n_fft.get() {
3493                return Err(SpectrogramError::invalid_input(format!(
3494                    "Custom window size ({}) must match n_fft ({})",
3495                    size.get(),
3496                    n_fft.get()
3497                )));
3498            }
3499        }
3500
3501        Ok(Self {
3502            n_fft,
3503            hop_size,
3504            window,
3505            centre,
3506        })
3507    }
3508
3509    const unsafe fn new_unchecked(
3510        n_fft: NonZeroUsize,
3511        hop_size: NonZeroUsize,
3512        window: WindowType,
3513        centre: bool,
3514    ) -> Self {
3515        Self {
3516            n_fft,
3517            hop_size,
3518            window,
3519            centre,
3520        }
3521    }
3522
3523    /// Get the FFT window size.
3524    ///
3525    /// # Returns
3526    ///
3527    /// The FFT window size.
3528    #[inline]
3529    #[must_use]
3530    pub const fn n_fft(&self) -> NonZeroUsize {
3531        self.n_fft
3532    }
3533
3534    /// Get the hop size (samples between successive frames).
3535    ///
3536    /// # Returns
3537    ///
3538    /// The hop size.
3539    #[inline]
3540    #[must_use]
3541    pub const fn hop_size(&self) -> NonZeroUsize {
3542        self.hop_size
3543    }
3544
3545    /// Get the window function.
3546    ///
3547    /// # Returns
3548    ///
3549    /// The window function.
3550    #[inline]
3551    #[must_use]
3552    pub fn window(&self) -> WindowType {
3553        self.window.clone()
3554    }
3555
3556    /// Get whether frames are centered (input signal is padded).
3557    ///
3558    /// # Returns
3559    ///
3560    /// `true` if frames are centered, `false` otherwise.
3561    #[inline]
3562    #[must_use]
3563    pub const fn centre(&self) -> bool {
3564        self.centre
3565    }
3566
3567    /// Create a builder for STFT parameters.
3568    ///
3569    /// # Returns
3570    ///
3571    /// A `StftParamsBuilder` instance.
3572    ///
3573    /// # Examples
3574    ///
3575    /// ```
3576    /// use spectrograms::{StftParams, WindowType, nzu};
3577    ///
3578    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
3579    /// let stft = StftParams::builder()
3580    ///     .n_fft(nzu!(2048))
3581    ///     .hop_size(nzu!(512))
3582    ///     .window(WindowType::Hanning)
3583    ///     .centre(true)
3584    ///     .build()?;
3585    ///
3586    /// assert_eq!(stft.n_fft(), nzu!(2048));
3587    /// assert_eq!(stft.hop_size(), nzu!(512));
3588    /// # Ok(())
3589    /// # }
3590    /// ```
3591    #[inline]
3592    #[must_use]
3593    pub fn builder() -> StftParamsBuilder {
3594        StftParamsBuilder::default()
3595    }
3596}
3597
3598/// Builder for [`StftParams`].
3599#[derive(Debug, Clone)]
3600pub struct StftParamsBuilder {
3601    n_fft: Option<NonZeroUsize>,
3602    hop_size: Option<NonZeroUsize>,
3603    window: WindowType,
3604    centre: bool,
3605}
3606
3607impl Default for StftParamsBuilder {
3608    #[inline]
3609    fn default() -> Self {
3610        Self {
3611            n_fft: None,
3612            hop_size: None,
3613            window: WindowType::Hanning,
3614            centre: true,
3615        }
3616    }
3617}
3618
3619impl StftParamsBuilder {
3620    /// Set the FFT window size.
3621    ///
3622    /// # Arguments
3623    ///
3624    /// * `n_fft` - Size of the FFT window
3625    ///
3626    /// # Returns
3627    ///
3628    /// The builder with the updated FFT window size.
3629    #[inline]
3630    #[must_use]
3631    pub const fn n_fft(mut self, n_fft: NonZeroUsize) -> Self {
3632        self.n_fft = Some(n_fft);
3633        self
3634    }
3635
3636    /// Set the hop size (samples between successive frames).
3637    ///
3638    /// # Arguments
3639    ///
3640    /// * `hop_size` - Number of samples between successive frames
3641    ///
3642    /// # Returns
3643    ///
3644    /// The builder with the updated hop size.
3645    #[inline]
3646    #[must_use]
3647    pub const fn hop_size(mut self, hop_size: NonZeroUsize) -> Self {
3648        self.hop_size = Some(hop_size);
3649        self
3650    }
3651
3652    /// Set the window function.
3653    ///
3654    /// # Arguments
3655    ///
3656    /// * `window` - Window function to apply to each frame
3657    ///
3658    /// # Returns
3659    ///
3660    /// The builder with the updated window function.
3661    #[inline]
3662    #[must_use]
3663    pub fn window(mut self, window: WindowType) -> Self {
3664        self.window = window;
3665        self
3666    }
3667
3668    /// Set whether to center frames (pad input signal).
3669    #[inline]
3670    #[must_use]
3671    pub const fn centre(mut self, centre: bool) -> Self {
3672        self.centre = centre;
3673        self
3674    }
3675
3676    /// Build the [`StftParams`].
3677    ///
3678    /// # Errors
3679    ///
3680    /// Returns an error if:
3681    /// - `n_fft` or `hop_size` are not set or are zero
3682    /// - `hop_size` > `n_fft`
3683    #[inline]
3684    pub fn build(self) -> SpectrogramResult<StftParams> {
3685        let n_fft = self
3686            .n_fft
3687            .ok_or_else(|| SpectrogramError::invalid_input("n_fft must be set"))?;
3688        let hop_size = self
3689            .hop_size
3690            .ok_or_else(|| SpectrogramError::invalid_input("hop_size must be set"))?;
3691
3692        StftParams::new(n_fft, hop_size, self.window, self.centre)
3693    }
3694}
3695
3696//
3697// ========================
3698// Mel parameters
3699// ========================
3700//
3701
3702/// Mel filterbank normalization strategy.
3703///
3704/// Determines how the triangular mel filters are normalized after construction.
3705#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3706#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3707#[non_exhaustive]
3708#[derive(Default)]
3709pub enum MelNorm {
3710    /// No normalization (triangular filters with peak = 1.0).
3711    ///
3712    /// This is the default and fastest option.
3713    #[default]
3714    None,
3715
3716    /// Slaney-style area normalization (librosa default).
3717    ///
3718    /// Each mel filter is divided by its bandwidth in Hz: `2.0 / (f_max - f_min)`.
3719    /// This ensures constant energy per mel band regardless of bandwidth.
3720    ///
3721    /// Use this for compatibility with librosa's default behavior.
3722    Slaney,
3723
3724    /// L1 normalization (sum of weights = 1.0).
3725    ///
3726    /// Each mel filter's weights are divided by their sum.
3727    /// Useful when you want each filter to act as a weighted average.
3728    L1,
3729
3730    /// L2 normalization (Euclidean norm = 1.0).
3731    ///
3732    /// Each mel filter's weights are divided by their L2 norm.
3733    /// Provides unit-norm filters in the L2 sense.
3734    L2,
3735}
3736
3737/// Mel filter bank parameters
3738///
3739/// * `n_mels`: Number of mel bands
3740/// * `f_min`: Minimum frequency (Hz)
3741/// * `f_max`: Maximum frequency (Hz)
3742/// * `norm`: Filterbank normalization strategy
3743#[derive(Debug, Clone, Copy, PartialEq)]
3744#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3745pub struct MelParams {
3746    n_mels: NonZeroUsize,
3747    f_min: f64,
3748    f_max: f64,
3749    norm: MelNorm,
3750}
3751
3752impl MelParams {
3753    /// Create new mel filter bank parameters.
3754    ///
3755    /// # Arguments
3756    ///
3757    /// * `n_mels` - Number of mel bands
3758    /// * `f_min` - Minimum frequency (Hz)
3759    /// * `f_max` - Maximum frequency (Hz)
3760    ///
3761    /// # Errors
3762    ///
3763    /// Returns an error if:
3764    /// - `f_min` is not >= 0
3765    /// - `f_max` is not > `f_min`
3766    ///
3767    /// # Returns
3768    ///
3769    /// A `MelParams` instance with no normalization (default).
3770    #[inline]
3771    pub fn new(n_mels: NonZeroUsize, f_min: f64, f_max: f64) -> SpectrogramResult<Self> {
3772        Self::with_norm(n_mels, f_min, f_max, MelNorm::None)
3773    }
3774
3775    /// Create new mel filter bank parameters with specified normalization.
3776    ///
3777    /// # Arguments
3778    ///
3779    /// * `n_mels` - Number of mel bands
3780    /// * `f_min` - Minimum frequency (Hz)
3781    /// * `f_max` - Maximum frequency (Hz)
3782    /// * `norm` - Filterbank normalization strategy
3783    ///
3784    /// # Errors
3785    ///
3786    /// Returns an error if:
3787    /// - `f_min` is not >= 0
3788    /// - `f_max` is not > `f_min`
3789    ///
3790    /// # Returns
3791    ///
3792    /// A `MelParams` instance.
3793    #[inline]
3794    pub fn with_norm(
3795        n_mels: NonZeroUsize,
3796        f_min: f64,
3797        f_max: f64,
3798        norm: MelNorm,
3799    ) -> SpectrogramResult<Self> {
3800        if f_min < 0.0 {
3801            return Err(SpectrogramError::invalid_input("f_min must be >= 0"));
3802        }
3803
3804        if f_max <= f_min {
3805            return Err(SpectrogramError::invalid_input("f_max must be > f_min"));
3806        }
3807
3808        Ok(Self {
3809            n_mels,
3810            f_min,
3811            f_max,
3812            norm,
3813        })
3814    }
3815
3816    /// Create new mel filter bank parameters.
3817    ///
3818    /// # Arguments
3819    ///
3820    /// * `n_mels` - Number of mel bands
3821    /// * `f_min` - Minimum frequency (Hz)
3822    /// * `f_max` - Maximum frequency (Hz)
3823    ///
3824    /// # Returns
3825    ///
3826    /// A `MelParams` instance with no normalization (default).
3827    ///
3828    /// # Safety
3829    ///
3830    /// The caller must ensure `f_min < f_max` (not validated at runtime).
3831    #[must_use]
3832    pub const unsafe fn new_unchecked(n_mels: NonZeroUsize, f_min: f64, f_max: f64) -> Self {
3833        Self {
3834            n_mels,
3835            f_min,
3836            f_max,
3837            norm: MelNorm::None,
3838        }
3839    }
3840
3841    /// Get the number of mel bands.
3842    ///
3843    /// # Returns
3844    ///
3845    /// The number of mel bands.
3846    #[inline]
3847    #[must_use]
3848    pub const fn n_mels(&self) -> NonZeroUsize {
3849        self.n_mels
3850    }
3851
3852    /// Get the minimum frequency (Hz).
3853    ///
3854    /// # Returns
3855    ///
3856    /// The minimum frequency in Hz.
3857    #[inline]
3858    #[must_use]
3859    pub const fn f_min(&self) -> f64 {
3860        self.f_min
3861    }
3862
3863    /// Get the maximum frequency (Hz).
3864    ///
3865    /// # Returns
3866    ///
3867    /// The maximum frequency in Hz.
3868    #[inline]
3869    #[must_use]
3870    pub const fn f_max(&self) -> f64 {
3871        self.f_max
3872    }
3873
3874    /// Get the filterbank normalization strategy.
3875    ///
3876    /// # Returns
3877    ///
3878    /// The normalization strategy.
3879    #[inline]
3880    #[must_use]
3881    pub const fn norm(&self) -> MelNorm {
3882        self.norm
3883    }
3884
3885    /// Create standard mel filterbank parameters.
3886    ///
3887    /// Uses 128 mel bands from 0 Hz to the Nyquist frequency.
3888    ///
3889    /// # Arguments
3890    ///
3891    /// * `sample_rate` - Sample rate in Hz (used to determine `f_max`)
3892    ///
3893    /// # Returns
3894    ///
3895    /// A `MelParams` instance with standard settings.
3896    ///
3897    /// # Panics
3898    ///
3899    /// Panics if `sample_rate` is not greater than 0.
3900    #[inline]
3901    #[must_use]
3902    pub const fn standard(sample_rate: f64) -> Self {
3903        assert!(sample_rate > 0.0);
3904        // safety: parameters are known to be valid
3905        unsafe { Self::new_unchecked(nzu!(128), 0.0, sample_rate / 2.0) }
3906    }
3907
3908    /// Create mel filterbank parameters optimized for speech.
3909    ///
3910    /// Uses 40 mel bands from 0 Hz to 8000 Hz (typical speech bandwidth).
3911    ///
3912    /// # Returns
3913    ///
3914    /// A `MelParams` instance with speech-optimized settings.
3915    #[inline]
3916    #[must_use]
3917    pub const fn speech_standard() -> Self {
3918        // safety: parameters are known to be valid
3919        unsafe { Self::new_unchecked(nzu!(40), 0.0, 8000.0) }
3920    }
3921}
3922
3923//
3924// ========================
3925// LogHz parameters
3926// ========================
3927//
3928
3929/// Logarithmic frequency scale parameters
3930///
3931/// * `n_bins`: Number of logarithmically-spaced frequency bins
3932/// * `f_min`: Minimum frequency (Hz)
3933/// * `f_max`: Maximum frequency (Hz)
3934#[derive(Debug, Clone, Copy, PartialEq)]
3935#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
3936pub struct LogHzParams {
3937    n_bins: NonZeroUsize,
3938    f_min: f64,
3939    f_max: f64,
3940}
3941
3942impl LogHzParams {
3943    /// Create new logarithmic frequency scale parameters.
3944    ///
3945    /// # Arguments
3946    ///
3947    /// * `n_bins` - Number of logarithmically-spaced frequency bins
3948    /// * `f_min` - Minimum frequency (Hz)
3949    /// * `f_max` - Maximum frequency (Hz)
3950    ///
3951    /// # Errors
3952    ///
3953    /// Returns an error if:
3954    /// - `f_min` is not finite and > 0
3955    /// - `f_max` is not > `f_min`
3956    ///
3957    /// # Returns
3958    ///
3959    /// A `LogHzParams` instance.
3960    #[inline]
3961    pub fn new(n_bins: NonZeroUsize, f_min: f64, f_max: f64) -> SpectrogramResult<Self> {
3962        if !(f_min > 0.0 && f_min.is_finite()) {
3963            return Err(SpectrogramError::invalid_input(
3964                "f_min must be finite and > 0",
3965            ));
3966        }
3967
3968        if f_max <= f_min {
3969            return Err(SpectrogramError::invalid_input("f_max must be > f_min"));
3970        }
3971
3972        Ok(Self {
3973            n_bins,
3974            f_min,
3975            f_max,
3976        })
3977    }
3978
3979    const unsafe fn new_unchecked(n_bins: NonZeroUsize, f_min: f64, f_max: f64) -> Self {
3980        Self {
3981            n_bins,
3982            f_min,
3983            f_max,
3984        }
3985    }
3986
3987    /// Get the number of frequency bins.
3988    ///
3989    /// # Returns
3990    ///
3991    /// The number of frequency bins.
3992    #[inline]
3993    #[must_use]
3994    pub const fn n_bins(&self) -> NonZeroUsize {
3995        self.n_bins
3996    }
3997
3998    /// Get the minimum frequency (Hz).
3999    ///
4000    /// # Returns
4001    ///
4002    /// The minimum frequency in Hz.
4003    #[inline]
4004    #[must_use]
4005    pub const fn f_min(&self) -> f64 {
4006        self.f_min
4007    }
4008
4009    /// Get the maximum frequency (Hz).
4010    ///
4011    /// # Returns
4012    ///
4013    /// The maximum frequency in Hz.
4014    #[inline]
4015    #[must_use]
4016    pub const fn f_max(&self) -> f64 {
4017        self.f_max
4018    }
4019
4020    /// Create standard logarithmic frequency parameters.
4021    ///
4022    /// Uses 128 log bins from 20 Hz to the Nyquist frequency.
4023    ///
4024    /// # Arguments
4025    ///
4026    /// * `sample_rate` - Sample rate in Hz (used to determine `f_max`)
4027    #[inline]
4028    #[must_use]
4029    pub fn standard(sample_rate: f64) -> Self {
4030        // safety: parameters are known to be valid
4031        unsafe { Self::new_unchecked(nzu!(128), 20.0, sample_rate / 2.0) }
4032    }
4033
4034    /// Create logarithmic frequency parameters optimized for music.
4035    ///
4036    /// Uses 84 bins (7 octaves * 12 bins/octave) from 27.5 Hz (A0) to 4186 Hz (C8).
4037    #[inline]
4038    #[must_use]
4039    pub const fn music_standard() -> Self {
4040        // safety: parameters are known to be valid
4041        unsafe { Self::new_unchecked(nzu!(84), 27.5, 4186.0) }
4042    }
4043}
4044
4045//
4046// ========================
4047// Log scaling parameters
4048// ========================
4049//
4050
4051#[derive(Debug, Clone, Copy, PartialEq)]
4052#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
4053pub struct LogParams {
4054    floor_db: f64,
4055}
4056
4057impl LogParams {
4058    /// Create new logarithmic scaling parameters.
4059    ///
4060    /// # Arguments
4061    ///
4062    /// * `floor_db` - Minimum dB value (floor) for logarithmic scaling
4063    ///
4064    /// # Errors
4065    ///
4066    /// Returns an error if `floor_db` is not finite.
4067    ///
4068    /// # Returns
4069    ///
4070    /// A `LogParams` instance.
4071    #[inline]
4072    pub fn new(floor_db: f64) -> SpectrogramResult<Self> {
4073        if !floor_db.is_finite() {
4074            return Err(SpectrogramError::invalid_input("floor_db must be finite"));
4075        }
4076
4077        Ok(Self { floor_db })
4078    }
4079
4080    /// Create new logarithmic scaling parameters.
4081    ///
4082    /// # Arguments
4083    ///
4084    /// * `floor_db` - Minimum dB value (floor) for logarithmic scaling
4085    ///
4086    /// # Returns
4087    ///
4088    /// A `LogParams` instance.
4089    #[inline]
4090    #[must_use]
4091    pub const fn new_unchecked(floor_db: f64) -> Self {
4092        Self { floor_db }
4093    }
4094
4095    /// Get the floor dB value.
4096    #[inline]
4097    #[must_use]
4098    pub const fn floor_db(&self) -> f64 {
4099        self.floor_db
4100    }
4101}
4102
4103/// Spectrogram computation parameters.
4104///
4105/// * `stft`: STFT parameters
4106/// * `sample_rate_hz`: Sample rate in Hz
4107#[derive(Debug, Clone)]
4108#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
4109pub struct SpectrogramParams {
4110    pub(crate) stft: StftParams,
4111    pub(crate) sample_rate_hz: f64,
4112}
4113
4114impl SpectrogramParams {
4115    /// Create new spectrogram parameters.
4116    ///
4117    /// # Arguments
4118    ///
4119    /// * `stft` - STFT parameters
4120    /// * `sample_rate_hz` - Sample rate in Hz
4121    ///
4122    /// # Errors
4123    ///
4124    /// Returns an error if the sample rate is not positive and finite.
4125    ///
4126    /// # Returns
4127    ///
4128    /// A `SpectrogramParams` instance.
4129    #[inline]
4130    pub fn new(stft: StftParams, sample_rate_hz: f64) -> SpectrogramResult<Self> {
4131        if !(sample_rate_hz > 0.0 && sample_rate_hz.is_finite()) {
4132            return Err(SpectrogramError::invalid_input(
4133                "sample_rate_hz must be finite and > 0",
4134            ));
4135        }
4136
4137        Ok(Self {
4138            stft,
4139            sample_rate_hz,
4140        })
4141    }
4142
4143    /// Create new spectrogram parameters without checking the arguments.
4144    ///
4145    /// # Arguments
4146    ///
4147    /// * `stft` - STFT parameters
4148    /// * `sample_rate_hz` - Sample rate in Hz
4149    ///
4150    /// # Returns
4151    ///
4152    /// A `SpectrogramParams` instance.
4153    #[inline]
4154    #[must_use]
4155    pub const fn new_unchecked(stft: StftParams, sample_rate_hz: f64) -> Self {
4156        Self {
4157            stft,
4158            sample_rate_hz,
4159        }
4160    }
4161
4162    /// Create a builder for spectrogram parameters.
4163    ///
4164    /// # Errors
4165    ///
4166    /// Returns an error if required parameters are not set or are invalid.
4167    ///
4168    /// # Returns
4169    ///
4170    /// A builder for [`SpectrogramParams`].
4171    ///
4172    /// # Examples
4173    ///
4174    /// ```
4175    /// use spectrograms::{SpectrogramParams, WindowType, nzu};
4176    ///
4177    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4178    /// let params = SpectrogramParams::builder()
4179    ///     .sample_rate(16000.0)
4180    ///     .n_fft(nzu!(512))
4181    ///     .hop_size(nzu!(256))
4182    ///     .window(WindowType::Hanning)
4183    ///     .centre(true)
4184    ///     .build()?;
4185    ///
4186    /// assert_eq!(params.sample_rate_hz(), 16000.0);
4187    /// # Ok(())
4188    /// # }
4189    /// ```
4190    #[inline]
4191    #[must_use]
4192    pub fn builder() -> SpectrogramParamsBuilder {
4193        SpectrogramParamsBuilder::default()
4194    }
4195
4196    /// Create default parameters for speech processing.
4197    ///
4198    /// # Arguments
4199    ///
4200    /// * `sample_rate_hz` - Sample rate in Hz
4201    ///
4202    /// # Returns
4203    ///
4204    /// A `SpectrogramParams` instance with default settings for music analysis.
4205    ///
4206    /// # Errors
4207    ///
4208    /// Returns an error if the sample rate is not positive and finite.
4209    ///
4210    /// Uses:
4211    /// - `n_fft`: 512 (32ms at 16kHz)
4212    /// - `hop_size`: 160 (10ms at 16kHz)
4213    /// - window: Hanning
4214    /// - centre: true
4215    #[inline]
4216    pub fn speech_default(sample_rate_hz: f64) -> SpectrogramResult<Self> {
4217        // safety: parameters are known to be valid
4218        let stft =
4219            unsafe { StftParams::new_unchecked(nzu!(512), nzu!(160), WindowType::Hanning, true) };
4220
4221        Self::new(stft, sample_rate_hz)
4222    }
4223
4224    /// Create default parameters for music processing.
4225    ///
4226    /// # Arguments
4227    ///
4228    /// * `sample_rate_hz` - Sample rate in Hz
4229    ///
4230    /// # Returns
4231    ///
4232    /// A `SpectrogramParams` instance with default settings for music analysis.
4233    ///
4234    /// # Errors
4235    ///
4236    /// Returns an error if the sample rate is not positive and finite.
4237    ///
4238    /// Uses:
4239    /// - `n_fft`: 2048 (46ms at 44.1kHz)
4240    /// - `hop_size`: 512 (11.6ms at 44.1kHz)
4241    /// - window: Hanning
4242    /// - centre: true
4243    #[inline]
4244    pub fn music_default(sample_rate_hz: f64) -> SpectrogramResult<Self> {
4245        // safety: parameters are known to be valid
4246        let stft =
4247            unsafe { StftParams::new_unchecked(nzu!(2048), nzu!(512), WindowType::Hanning, true) };
4248        Self::new(stft, sample_rate_hz)
4249    }
4250
4251    /// Get the STFT parameters.
4252    #[inline]
4253    #[must_use]
4254    pub const fn stft(&self) -> &StftParams {
4255        &self.stft
4256    }
4257
4258    /// Get the sample rate in Hz.
4259    #[inline]
4260    #[must_use]
4261    pub const fn sample_rate_hz(&self) -> f64 {
4262        self.sample_rate_hz
4263    }
4264
4265    /// Get the frame period in seconds.
4266    #[inline]
4267    #[must_use]
4268    #[allow(clippy::cast_precision_loss)]
4269    pub fn frame_period_seconds(&self) -> f64 {
4270        self.stft.hop_size().get() as f64 / self.sample_rate_hz
4271    }
4272
4273    /// Get the Nyquist frequency in Hz.
4274    #[inline]
4275    #[must_use]
4276    pub fn nyquist_hz(&self) -> f64 {
4277        self.sample_rate_hz * 0.5
4278    }
4279}
4280
4281/// Builder for [`SpectrogramParams`].
4282#[derive(Debug, Clone)]
4283pub struct SpectrogramParamsBuilder {
4284    sample_rate: Option<f64>,
4285    n_fft: Option<NonZeroUsize>,
4286    hop_size: Option<NonZeroUsize>,
4287    window: WindowType,
4288    centre: bool,
4289}
4290
4291impl Default for SpectrogramParamsBuilder {
4292    #[inline]
4293    fn default() -> Self {
4294        Self {
4295            sample_rate: None,
4296            n_fft: None,
4297            hop_size: None,
4298            window: WindowType::Hanning,
4299            centre: true,
4300        }
4301    }
4302}
4303
4304impl SpectrogramParamsBuilder {
4305    /// Set the sample rate in Hz.
4306    ///
4307    /// # Arguments
4308    ///
4309    /// * `sample_rate` - Sample rate in Hz.
4310    ///
4311    /// # Returns
4312    ///
4313    /// The updated builder instance.
4314    #[inline]
4315    #[must_use]
4316    pub const fn sample_rate(mut self, sample_rate: f64) -> Self {
4317        self.sample_rate = Some(sample_rate);
4318        self
4319    }
4320
4321    /// Set the FFT window size.
4322    ///
4323    /// # Arguments
4324    ///
4325    /// * `n_fft` - FFT size.
4326    ///
4327    /// # Returns
4328    ///
4329    /// The updated builder instance.
4330    #[inline]
4331    #[must_use]
4332    pub const fn n_fft(mut self, n_fft: NonZeroUsize) -> Self {
4333        self.n_fft = Some(n_fft);
4334        self
4335    }
4336
4337    /// Set the hop size (samples between successive frames).
4338    ///
4339    /// # Arguments
4340    ///
4341    /// * `hop_size` - Hop size in samples.
4342    ///
4343    /// # Returns
4344    ///
4345    /// The updated builder instance.
4346    #[inline]
4347    #[must_use]
4348    pub const fn hop_size(mut self, hop_size: NonZeroUsize) -> Self {
4349        self.hop_size = Some(hop_size);
4350        self
4351    }
4352
4353    /// Set the window function.
4354    ///
4355    /// # Arguments
4356    ///
4357    /// * `window` - Window function to apply to each frame.
4358    ///
4359    /// # Returns
4360    ///
4361    /// The updated builder instance.
4362    #[inline]
4363    #[must_use]
4364    pub fn window(mut self, window: WindowType) -> Self {
4365        self.window = window;
4366        self
4367    }
4368
4369    /// Set whether to center frames (pad input signal).
4370    ///
4371    /// # Arguments
4372    ///
4373    /// * `centre` - If true, frames are centered by padding the input signal.
4374    ///
4375    /// # Returns
4376    ///
4377    /// The updated builder instance.
4378    #[inline]
4379    #[must_use]
4380    pub const fn centre(mut self, centre: bool) -> Self {
4381        self.centre = centre;
4382        self
4383    }
4384
4385    /// Build the [`SpectrogramParams`].
4386    ///
4387    /// # Errors
4388    ///
4389    /// Returns an error if required parameters are not set or are invalid.
4390    ///
4391    /// # Returns
4392    ///
4393    /// A `SpectrogramParams` instance.
4394    #[inline]
4395    pub fn build(self) -> SpectrogramResult<SpectrogramParams> {
4396        let sample_rate = self
4397            .sample_rate
4398            .ok_or_else(|| SpectrogramError::invalid_input("sample_rate must be set"))?;
4399        let n_fft = self
4400            .n_fft
4401            .ok_or_else(|| SpectrogramError::invalid_input("n_fft must be set"))?;
4402        let hop_size = self
4403            .hop_size
4404            .ok_or_else(|| SpectrogramError::invalid_input("hop_size must be set"))?;
4405
4406        let stft = StftParams::new(n_fft, hop_size, self.window, self.centre)?;
4407        SpectrogramParams::new(stft, sample_rate)
4408    }
4409
4410    /// Build the [`SpectrogramParams`].
4411    ///
4412    /// # Safety
4413    ///
4414    /// The caller is responsible for ensuring the 'n_fft', 'hop_size' are set.
4415    ///
4416    /// # Returns
4417    ///
4418    /// A `SpectrogramParams` instance.
4419    #[inline]
4420    #[must_use]
4421    pub unsafe fn build_unchecked(self) -> SpectrogramParams {
4422        // safety: is the repsonsibility of the caller
4423        unsafe {
4424            let n_fft = self.n_fft.unwrap_unchecked();
4425            let hop_size = self.hop_size.unwrap_unchecked();
4426            let stft = StftParams::new_unchecked(n_fft, hop_size, self.window, self.centre);
4427            let sample_rate = self.sample_rate.unwrap_unchecked();
4428            SpectrogramParams::new_unchecked(stft, sample_rate)
4429        }
4430    }
4431}
4432
4433//
4434// ========================
4435// Standalone FFT Functions
4436// ========================
4437//
4438
4439/// Compute the real-to-complex FFT of a real-valued signal.
4440///
4441/// This function performs a forward FFT on real-valued input, returning the
4442/// complex frequency domain representation. Only the positive frequencies
4443/// are returned (length = `n_fft/2` + 1) due to conjugate symmetry.
4444///
4445/// # Arguments
4446///
4447/// * `samples` - Input signal (length ≤ n_fft, will be zero-padded if shorter)
4448/// * `n_fft` - FFT size
4449///
4450/// # Returns
4451///
4452/// A vector of complex frequency bins with length `n_fft/2` + 1.
4453///
4454/// # Automatic Zero-Padding
4455///
4456/// If the input signal is shorter than `n_fft`, it will be automatically
4457/// zero-padded to the required length. This is standard DSP practice and
4458/// preserves frequency resolution (bin spacing = sample_rate / n_fft).
4459///
4460/// ```
4461/// use spectrograms::{fft, nzu};
4462/// use non_empty_slice::non_empty_vec;
4463///
4464/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4465/// let signal = non_empty_vec![1.0, 2.0, 3.0]; // Only 3 samples
4466/// let spectrum = fft(&signal, nzu!(8))?;   // Automatically padded to 8
4467/// assert_eq!(spectrum.len(), 5);     // Output: 8/2 + 1 = 5 bins
4468/// # Ok(())
4469/// # }
4470/// ```
4471///
4472/// # Errors
4473///
4474/// Returns `InvalidInput` error if the input length exceeds `n_fft`.
4475///
4476/// # Examples
4477///
4478/// ```
4479/// use spectrograms::*;
4480/// use non_empty_slice::non_empty_vec;
4481///
4482/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4483/// let signal = non_empty_vec![0.0; nzu!(512)];
4484/// let spectrum = fft(&signal, nzu!(512))?;
4485///
4486/// assert_eq!(spectrum.len(), 257); // 512/2 + 1
4487/// # Ok(())
4488/// # }
4489/// ```
4490#[inline]
4491pub fn fft<T: Sample>(
4492    samples: &NonEmptySlice<T>,
4493    n_fft: NonZeroUsize,
4494) -> SpectrogramResult<Array1<Complex<T>>> {
4495    if samples.len() > n_fft {
4496        return Err(SpectrogramError::invalid_input(format!(
4497            "Input length ({}) exceeds FFT size ({})",
4498            samples.len(),
4499            n_fft
4500        )));
4501    }
4502
4503    let out_len = r2c_output_size(n_fft.get());
4504
4505    let mut fft = T::plan_r2c(n_fft.get())?;
4506
4507    let input = if samples.len() < n_fft {
4508        let mut padded = vec![T::zero(); n_fft.get()];
4509        padded[..samples.len().get()].copy_from_slice(samples);
4510        // safety: samples.len() < n_fft checked above and n_fft > 0
4511
4512        unsafe { NonEmptyVec::new_unchecked(padded) }
4513    } else {
4514        samples.to_non_empty_vec()
4515    };
4516
4517    let mut output = vec![Complex::new(T::zero(), T::zero()); out_len];
4518    fft.process(&input, &mut output)?;
4519    let output = Array1::from_vec(output);
4520    Ok(output)
4521}
4522
4523#[inline]
4524/// Compute the real-valued fft of a signal.
4525///
4526/// # Arguments
4527/// * `samples` - Input signal (length ≤ n_fft, will be zero-padded if shorter)
4528/// * `n_fft` - FFT size
4529///
4530/// # Returns
4531///
4532/// An array with length `n_fft/2` + 1.
4533///
4534/// # Errors
4535///
4536/// Returns `InvalidInput` error if the input length exceeds `n_fft`.
4537///
4538/// # Examples
4539///
4540/// ```
4541/// use spectrograms::*;
4542/// use non_empty_slice::non_empty_vec;
4543///
4544/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4545/// let signal = non_empty_vec![0.0; nzu!(512)];
4546/// let rfft_result = rfft(&signal, nzu!(512))?;
4547/// // equivalent to
4548/// let fft_result = fft(&signal, nzu!(512))?;
4549/// let rfft_result = fft_result.mapv(num_complex::Complex::norm);
4550/// # Ok(())
4551/// # }
4552///
4553pub fn rfft<T: Sample>(
4554    samples: &NonEmptySlice<T>,
4555    n_fft: NonZeroUsize,
4556) -> SpectrogramResult<Array1<T>> {
4557    Ok(fft(samples, n_fft)?.mapv(Complex::norm))
4558}
4559
4560/// Compute the power spectrum of a signal (|X|²).
4561///
4562/// This function applies an optional window function and computes the
4563/// power spectrum via FFT. The result contains only positive frequencies.
4564///
4565/// # Arguments
4566///
4567/// * `samples` - Input signal (length ≤ n_fft, will be zero-padded if shorter)
4568/// * `n_fft` - FFT size
4569/// * `window` - Optional window function (None for rectangular window)
4570///
4571/// # Returns
4572///
4573/// A vector of power values with length `n_fft/2` + 1.
4574///
4575/// # Automatic Zero-Padding
4576///
4577/// If the input signal is shorter than `n_fft`, it will be automatically
4578/// zero-padded to the required length. This is standard DSP practice and
4579/// preserves frequency resolution (bin spacing = sample_rate / n_fft).
4580///
4581/// ```
4582/// use spectrograms::{power_spectrum, nzu};
4583/// use non_empty_slice::non_empty_vec;
4584///
4585/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4586/// let signal = non_empty_vec![1.0, 2.0, 3.0]; // Only 3 samples
4587/// let power = power_spectrum(&signal, nzu!(8), None)?;
4588/// assert_eq!(power.len(), nzu!(5));     // Output: 8/2 + 1 = 5 bins
4589/// # Ok(())
4590/// # }
4591/// ```
4592///
4593/// # Errors
4594///
4595/// Returns `InvalidInput` error if the input length exceeds `n_fft`.
4596///
4597/// # Examples
4598///
4599/// ```
4600/// use spectrograms::*;
4601/// use non_empty_slice::non_empty_vec;
4602///
4603/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4604/// let signal = non_empty_vec![0.0; nzu!(512)];
4605/// let power = power_spectrum(&signal, nzu!(512), Some(WindowType::Hanning))?;
4606///
4607/// assert_eq!(power.len(), nzu!(257)); // 512/2 + 1
4608/// # Ok(())
4609/// # }
4610/// ```
4611#[inline]
4612pub fn power_spectrum<T: Sample>(
4613    samples: &NonEmptySlice<T>,
4614    n_fft: NonZeroUsize,
4615    window: Option<WindowType>,
4616) -> SpectrogramResult<NonEmptyVec<T>> {
4617    if samples.len() > n_fft {
4618        return Err(SpectrogramError::invalid_input(format!(
4619            "Input length ({}) exceeds FFT size ({})",
4620            samples.len(),
4621            n_fft
4622        )));
4623    }
4624
4625    let mut windowed = vec![T::zero(); n_fft.get()];
4626    windowed[..samples.len().get()].copy_from_slice(samples);
4627
4628    if let Some(win_type) = window {
4629        let window_samples = make_window::<T>(win_type, n_fft);
4630        for i in 0..n_fft.get() {
4631            windowed[i] *= window_samples[i];
4632        }
4633    }
4634
4635    // safety: windowed is non-empty since n_fft > 0
4636    let windowed = unsafe { NonEmptySlice::new_unchecked(&windowed) };
4637    let fft_result = fft(windowed, n_fft)?;
4638    let fft_result: Vec<T> = fft_result
4639        .iter()
4640        .map(num_complex::Complex::norm_sqr)
4641        .collect();
4642    // safety: fft_result is non-empty since fft returned successfully
4643    Ok(unsafe { NonEmptyVec::new_unchecked(fft_result) })
4644}
4645
4646/// Compute the magnitude spectrum of a signal (|X|).
4647///
4648/// This function applies an optional window function and computes the
4649/// magnitude spectrum via FFT. The result contains only positive frequencies.
4650///
4651/// # Arguments
4652///
4653/// * `samples` - Input signal (length ≤ n_fft, will be zero-padded if shorter)
4654/// * `n_fft` - FFT size
4655/// * `window` - Optional window function (None for rectangular window)
4656///
4657/// # Automatic Zero-Padding
4658///
4659/// If the input signal is shorter than `n_fft`, it will be automatically
4660/// zero-padded to the required length. This preserves frequency resolution.
4661///
4662/// # Errors
4663///
4664/// Returns `InvalidInput` error if the input length exceeds `n_fft`.
4665///
4666/// # Returns
4667///
4668/// A vector of magnitude values with length `n_fft/2` + 1.
4669///
4670/// # Examples
4671///
4672/// ```
4673/// use spectrograms::*;
4674/// use non_empty_slice::non_empty_vec;
4675///
4676/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4677/// let signal = non_empty_vec![0.0; nzu!(512)];
4678/// let magnitude = magnitude_spectrum(&signal, nzu!(512), Some(WindowType::Hanning))?;
4679///
4680/// assert_eq!(magnitude.len(), nzu!(257)); // 512/2 + 1
4681/// # Ok(())
4682/// # }
4683/// ```
4684#[inline]
4685pub fn magnitude_spectrum<T: Sample>(
4686    samples: &NonEmptySlice<T>,
4687    n_fft: NonZeroUsize,
4688    window: Option<WindowType>,
4689) -> SpectrogramResult<NonEmptyVec<T>> {
4690    let power = power_spectrum(samples, n_fft, window)?;
4691    let power: Vec<T> = power.iter().map(|&p| p.sqrt()).collect();
4692    // safety: power is non-empty since power_spectrum returned successfully
4693    Ok(unsafe { NonEmptyVec::new_unchecked(power) })
4694}
4695
4696/// Compute the Short-Time Fourier Transform (STFT) of a signal.
4697///
4698/// This function computes the STFT by applying a sliding window and FFT
4699/// to sequential frames of the input signal.
4700///
4701/// # Arguments
4702///
4703/// * `samples` - Input signal (any type that can be converted to a slice)
4704/// * `n_fft` - FFT size
4705/// * `hop_size` - Number of samples between successive frames
4706/// * `window` - Window function to apply to each frame
4707/// * `center` - If true, pad the signal to center frames
4708///
4709/// # Returns
4710///
4711/// A 2D array with shape (`frequency_bins`, `time_frames`) containing complex STFT values.
4712///
4713/// # Errors
4714///
4715/// Returns an error if:
4716/// - `hop_size` > `n_fft`
4717/// - STFT computation fails
4718///
4719/// # Examples
4720///
4721/// ```
4722/// use spectrograms::*;
4723/// use non_empty_slice::non_empty_vec;
4724///
4725/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4726/// let signal = non_empty_vec![0.0; nzu!(16000)];
4727/// let stft_matrix = stft(&signal, nzu!(512), nzu!(256), WindowType::Hanning, true)?;
4728///
4729/// println!("STFT: {} bins x {} frames", stft_matrix.nrows(), stft_matrix.ncols());
4730/// # Ok(())
4731/// # }
4732/// ```
4733#[inline]
4734pub fn stft<T: Sample>(
4735    samples: &NonEmptySlice<T>,
4736    n_fft: NonZeroUsize,
4737    hop_size: NonZeroUsize,
4738    window: WindowType,
4739    center: bool,
4740) -> SpectrogramResult<Array2<Complex<T>>> {
4741    let stft_params = StftParams::new(n_fft, hop_size, window, center)?;
4742    let params = SpectrogramParams::new(stft_params, 1.0)?; // dummy sample rate
4743
4744    let mut plan = StftPlan::<T>::new(&params)?;
4745    let result = plan.compute(samples, &params)?;
4746
4747    Ok(result.data)
4748}
4749
4750/// Compute the inverse real FFT (complex-to-real IFFT).
4751///
4752/// This function performs an inverse FFT, converting frequency domain data
4753/// back to the time domain. Only the positive frequencies need to be provided
4754/// (length = `n_fft/2` + 1) due to conjugate symmetry.
4755///
4756/// # Arguments
4757///
4758/// * `spectrum` - Complex frequency bins (length should be `n_fft/2` + 1)
4759/// * `n_fft` - FFT size (length of the output signal)
4760///
4761/// # Returns
4762///
4763/// A vector of real-valued time-domain samples with length `n_fft`.
4764///
4765/// # Errors
4766///
4767/// Returns an error if:
4768/// - `spectrum` length doesn't match `n_fft/2` + 1
4769/// - Inverse FFT computation fails
4770///
4771/// # Examples
4772///
4773/// ```
4774/// use spectrograms::*;
4775/// use non_empty_slice::{non_empty_vec, NonEmptySlice};
4776/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4777/// // Forward FFT
4778/// let signal = non_empty_vec![1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0, 0.0];
4779/// let spectrum = fft(&signal, nzu!(8))?;
4780/// let slice = spectrum.as_slice().unwrap();
4781/// let spectrum_slice = NonEmptySlice::new(slice).unwrap();
4782/// // Inverse FFT
4783/// let reconstructed = irfft(spectrum_slice, nzu!(8))?;
4784///
4785/// assert_eq!(reconstructed.len(), nzu!(8));
4786/// # Ok(())
4787/// # }
4788/// ```
4789#[inline]
4790pub fn irfft<T: Sample>(
4791    spectrum: &NonEmptySlice<Complex<T>>,
4792    n_fft: NonZeroUsize,
4793) -> SpectrogramResult<NonEmptyVec<T>> {
4794    use crate::fft_backend::{C2rPlan, r2c_output_size};
4795
4796    let n_fft = n_fft.get();
4797    let expected_len = r2c_output_size(n_fft);
4798    if spectrum.len().get() != expected_len {
4799        return Err(SpectrogramError::dimension_mismatch(
4800            expected_len,
4801            spectrum.len().get(),
4802        ));
4803    }
4804
4805    let mut ifft = T::plan_c2r(n_fft)?;
4806
4807    let mut output = vec![T::zero(); n_fft];
4808    ifft.process(spectrum.as_slice(), &mut output)?;
4809
4810    // Safety: output is non-empty since n_fft > 0
4811    Ok(unsafe { NonEmptyVec::new_unchecked(output) })
4812}
4813
4814/// Reconstruct a time-domain signal from its STFT using overlap-add.
4815///
4816/// This function performs the inverse Short-Time Fourier Transform, converting
4817/// a 2D complex STFT matrix back to a 1D time-domain signal using overlap-add
4818/// synthesis with the specified window function.
4819///
4820/// # Arguments
4821///
4822/// * `stft_matrix` - Complex STFT with shape (`frequency_bins`, `time_frames`)
4823/// * `n_fft` - FFT size
4824/// * `hop_size` - Number of samples between successive frames
4825/// * `window` - Window function to apply (should match forward STFT window)
4826/// * `center` - If true, assume the forward STFT was centered
4827///
4828/// # Returns
4829///
4830/// A vector of reconstructed time-domain samples.
4831///
4832/// # Errors
4833///
4834/// Returns an error if:
4835/// - `stft_matrix` dimensions are inconsistent with `n_fft`
4836/// - `hop_size` > `n_fft`
4837/// - Inverse STFT computation fails
4838///
4839/// # Examples
4840///
4841/// ```
4842/// use spectrograms::*;
4843/// use non_empty_slice::non_empty_vec;
4844///
4845/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4846/// // Generate signal
4847/// let signal = non_empty_vec![1.0; nzu!(16000)];
4848///
4849/// // Forward STFT
4850/// let stft_matrix = stft(&signal, nzu!(512), nzu!(256), WindowType::Hanning, true)?;
4851///
4852/// // Inverse STFT
4853/// let reconstructed = istft(&stft_matrix, nzu!(512), nzu!(256), WindowType::Hanning, true)?;
4854///
4855/// println!("Original: {} samples", signal.len());
4856/// println!("Reconstructed: {} samples", reconstructed.len());
4857/// # Ok(())
4858/// # }
4859/// ```
4860#[inline]
4861pub fn istft<T: Sample>(
4862    stft_matrix: &Array2<Complex<T>>,
4863    n_fft: NonZeroUsize,
4864    hop_size: NonZeroUsize,
4865    window: WindowType,
4866    center: bool,
4867) -> SpectrogramResult<NonEmptyVec<T>> {
4868    use crate::fft_backend::{C2rPlan, r2c_output_size};
4869
4870    let n_bins = stft_matrix.nrows();
4871    let n_frames = stft_matrix.ncols();
4872
4873    let expected_bins = r2c_output_size(n_fft.get());
4874    if n_bins != expected_bins {
4875        return Err(SpectrogramError::dimension_mismatch(expected_bins, n_bins));
4876    }
4877    if hop_size.get() > n_fft.get() {
4878        return Err(SpectrogramError::invalid_input("hop_size must be <= n_fft"));
4879    }
4880    // Create inverse FFT plan
4881    let mut ifft = T::plan_c2r(n_fft.get())?;
4882
4883    // Generate window
4884    let window_samples = make_window::<T>(window, n_fft);
4885    let n_fft = n_fft.get();
4886    let hop_size = hop_size.get();
4887    // Calculate output length
4888    let pad = if center { n_fft / 2 } else { 0 };
4889    let output_len = (n_frames - 1) * hop_size + n_fft;
4890    // safety: output_len > 0 since n_frames > 0 and n_fft, hop_size > 0
4891    let output_len = unsafe { NonZeroUsize::new_unchecked(output_len) };
4892    let unpadded_len = output_len.get().saturating_sub(2 * pad);
4893
4894    // Allocate output buffer and normalization buffer
4895    let mut output = non_empty_vec![T::zero(); output_len];
4896    let mut norm = non_empty_vec![T::zero(); output_len];
4897
4898    // Overlap-add synthesis
4899    let mut frame_buffer = vec![Complex::new(T::zero(), T::zero()); n_bins];
4900    let mut time_frame = vec![T::zero(); n_fft];
4901
4902    for frame_idx in 0..n_frames {
4903        // Extract complex frame from STFT matrix
4904        for bin_idx in 0..n_bins {
4905            frame_buffer[bin_idx] = stft_matrix[[bin_idx, frame_idx]];
4906        }
4907
4908        // Inverse FFT
4909        ifft.process(&frame_buffer, &mut time_frame)?;
4910
4911        // Apply window
4912        for i in 0..n_fft {
4913            time_frame[i] *= window_samples[i];
4914        }
4915
4916        // Overlap-add into output buffer
4917        let start = frame_idx * hop_size;
4918        for i in 0..n_fft {
4919            let pos = start + i;
4920            if pos < output_len.get() {
4921                output[pos] += time_frame[i];
4922                norm[pos] += window_samples[i] * window_samples[i];
4923            }
4924        }
4925    }
4926
4927    // Normalize by window energy
4928    let norm_eps = T::from_f64(1e-10);
4929    for i in 0..output_len.get() {
4930        if norm[i] > norm_eps {
4931            output[i] /= norm[i];
4932        }
4933    }
4934
4935    // Remove padding if centered
4936    if center && unpadded_len > 0 {
4937        let start = pad;
4938        let end = start + unpadded_len;
4939        // safety: start < end <= output_len, therefore slice is non-empty
4940        output = unsafe {
4941            NonEmptySlice::new_unchecked(&output[start..end.min(output_len.get())])
4942                .to_non_empty_vec()
4943        };
4944    }
4945
4946    Ok(output)
4947}
4948
4949//
4950// ========================
4951// Reusable FFT Plans
4952// ========================
4953//
4954
4955/// A reusable FFT planner for efficient repeated FFT operations.
4956///
4957/// This planner caches FFT plans internally, making repeated FFT operations
4958/// of the same size much more efficient than calling `fft()` repeatedly.
4959///
4960/// # Examples
4961///
4962/// ```
4963/// use spectrograms::*;
4964/// use non_empty_slice::non_empty_vec;
4965///
4966/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
4967/// let mut planner = FftPlanner::new();
4968///
4969/// // Process multiple signals of the same size efficiently
4970/// for _ in 0..100 {
4971///     let signal = non_empty_vec![0.0; nzu!(512)];
4972///     let spectrum = planner.fft(&signal, nzu!(512))?;
4973///     // ... process spectrum ...
4974/// }
4975/// # Ok(())
4976/// # }
4977/// ```
4978pub struct FftPlanner {
4979    #[cfg(feature = "realfft")]
4980    inner: crate::RealFftPlanner,
4981    #[cfg(feature = "fftw")]
4982    inner: crate::FftwPlanner,
4983}
4984
4985impl FftPlanner {
4986    /// Create a new FFT planner with empty cache.
4987    #[inline]
4988    #[must_use]
4989    pub fn new() -> Self {
4990        Self {
4991            #[cfg(feature = "realfft")]
4992            inner: crate::RealFftPlanner::new(),
4993            #[cfg(feature = "fftw")]
4994            inner: crate::FftwPlanner::new(),
4995        }
4996    }
4997
4998    /// Compute forward FFT, reusing cached plans.
4999    ///
5000    /// This is more efficient than calling the standalone `fft()` function
5001    /// repeatedly for the same FFT size.
5002    ///
5003    /// # Automatic Zero-Padding
5004    ///
5005    /// If the input signal is shorter than `n_fft`, it will be automatically
5006    /// zero-padded to the required length.
5007    ///
5008    /// # Errors
5009    ///
5010    /// Returns `InvalidInput` error if the input length exceeds `n_fft`.
5011    ///
5012    /// # Examples
5013    ///
5014    /// ```
5015    /// use spectrograms::*;
5016    /// use non_empty_slice::non_empty_vec;
5017    ///
5018    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
5019    /// let mut planner = FftPlanner::new();
5020    ///
5021    /// let signal = non_empty_vec![1.0; nzu!(512)];
5022    /// let spectrum = planner.fft(&signal, nzu!(512))?;
5023    ///
5024    /// assert_eq!(spectrum.len(), 257); // 512/2 + 1
5025    /// # Ok(())
5026    /// # }
5027    /// ```
5028    #[inline]
5029    pub fn fft(
5030        &mut self,
5031        samples: &NonEmptySlice<f64>,
5032        n_fft: NonZeroUsize,
5033    ) -> SpectrogramResult<Array1<Complex<f64>>> {
5034        use crate::fft_backend::{R2cPlan, R2cPlanner, r2c_output_size};
5035
5036        if samples.len() > n_fft {
5037            return Err(SpectrogramError::invalid_input(format!(
5038                "Input length ({}) exceeds FFT size ({})",
5039                samples.len(),
5040                n_fft
5041            )));
5042        }
5043
5044        let out_len = r2c_output_size(n_fft.get());
5045        let mut plan = self.inner.plan_r2c(n_fft.get())?;
5046
5047        let input = if samples.len() < n_fft {
5048            let mut padded = vec![0.0; n_fft.get()];
5049            padded[..samples.len().get()].copy_from_slice(samples);
5050
5051            // safety: samples.len() < n_fft checked above and n_fft > 0
5052            unsafe { NonEmptyVec::new_unchecked(padded) }
5053        } else {
5054            samples.to_non_empty_vec()
5055        };
5056
5057        let mut output = vec![Complex::new(0.0, 0.0); out_len];
5058        plan.process(&input, &mut output)?;
5059
5060        let output = Array1::from_vec(output);
5061        Ok(output)
5062    }
5063
5064    /// Compute forward real FFT magnitude
5065    ///
5066    /// # Errors
5067    ///
5068    /// Returns an error if:
5069    /// - `n_fft` doesn't match the samples length
5070    ///
5071    ///
5072    #[inline]
5073    pub fn rfft(
5074        &mut self,
5075        samples: &NonEmptySlice<f64>,
5076        n_fft: NonZeroUsize,
5077    ) -> SpectrogramResult<Array1<f64>> {
5078        let fft_with_complex = fft(samples, n_fft)?;
5079        Ok(fft_with_complex.mapv(Complex::norm))
5080    }
5081
5082    /// Compute inverse FFT, reusing cached plans.
5083    ///
5084    /// This is more efficient than calling the standalone `irfft()` function
5085    /// repeatedly for the same FFT size.
5086    ///
5087    /// # Errors
5088    /// Returns an error if:
5089    ///
5090    /// - The calculated expected length of `spectrum` doesn't match its actual length
5091    ///
5092    /// # Examples
5093    ///
5094    /// ```
5095    /// use spectrograms::*;
5096    /// use non_empty_slice::{non_empty_vec, NonEmptySlice};
5097    ///
5098    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
5099    /// let mut planner = FftPlanner::new();
5100    ///
5101    /// // Forward FFT
5102    /// let signal = non_empty_vec![1.0; nzu!(512)];
5103    /// let spectrum = planner.fft(&signal, nzu!(512))?;
5104    ///
5105    /// // Inverse FFT
5106    /// let spectrum_slice = NonEmptySlice::new(spectrum.as_slice().unwrap()).unwrap();
5107    /// let reconstructed = planner.irfft(spectrum_slice, nzu!(512))?;
5108    ///
5109    /// assert_eq!(reconstructed.len(), nzu!(512));
5110    /// # Ok(())
5111    /// # }
5112    /// ```
5113    #[inline]
5114    pub fn irfft(
5115        &mut self,
5116        spectrum: &NonEmptySlice<Complex<f64>>,
5117        n_fft: NonZeroUsize,
5118    ) -> SpectrogramResult<NonEmptyVec<f64>> {
5119        use crate::fft_backend::{C2rPlan, C2rPlanner, r2c_output_size};
5120
5121        let expected_len = r2c_output_size(n_fft.get());
5122        if spectrum.len().get() != expected_len {
5123            return Err(SpectrogramError::dimension_mismatch(
5124                expected_len,
5125                spectrum.len().get(),
5126            ));
5127        }
5128
5129        let mut plan = self.inner.plan_c2r(n_fft.get())?;
5130        let mut output = vec![0.0; n_fft.get()];
5131        plan.process(spectrum, &mut output)?;
5132        // Safety: output is non-empty since n_fft > 0
5133        let output = unsafe { NonEmptyVec::new_unchecked(output) };
5134        Ok(output)
5135    }
5136
5137    /// Compute power spectrum with optional windowing, reusing cached plans.
5138    ///
5139    /// # Automatic Zero-Padding
5140    ///
5141    /// If the input signal is shorter than `n_fft`, it will be automatically
5142    /// zero-padded to the required length.
5143    ///
5144    /// # Errors
5145    /// Returns `InvalidInput` error if the input length exceeds `n_fft`.
5146    ///
5147    /// # Examples
5148    ///
5149    /// ```
5150    /// use spectrograms::*;
5151    /// use non_empty_slice::non_empty_vec;
5152    ///
5153    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
5154    /// let mut planner = FftPlanner::new();
5155    ///
5156    /// let signal = non_empty_vec![1.0; nzu!(512)];
5157    /// let power = planner.power_spectrum(&signal, nzu!(512), Some(WindowType::Hanning))?;
5158    ///
5159    /// assert_eq!(power.len(), nzu!(257));
5160    /// # Ok(())
5161    /// # }
5162    /// ```
5163    #[inline]
5164    pub fn power_spectrum(
5165        &mut self,
5166        samples: &NonEmptySlice<f64>,
5167        n_fft: NonZeroUsize,
5168        window: Option<WindowType>,
5169    ) -> SpectrogramResult<NonEmptyVec<f64>> {
5170        if samples.len() > n_fft {
5171            return Err(SpectrogramError::invalid_input(format!(
5172                "Input length ({}) exceeds FFT size ({})",
5173                samples.len(),
5174                n_fft
5175            )));
5176        }
5177
5178        let mut windowed = vec![0.0; n_fft.get()];
5179        windowed[..samples.len().get()].copy_from_slice(samples);
5180        if let Some(win_type) = window {
5181            let window_samples = make_window::<f64>(win_type, n_fft);
5182            for i in 0..n_fft.get() {
5183                windowed[i] *= window_samples[i];
5184            }
5185        }
5186
5187        // safety: windowed is non-empty since n_fft > 0
5188        let windowed = unsafe { NonEmptySlice::new_unchecked(&windowed) };
5189        let fft_result = self.fft(windowed, n_fft)?;
5190        let f = fft_result
5191            .iter()
5192            .map(num_complex::Complex::norm_sqr)
5193            .collect();
5194        // safety: fft_result is non-empty since fft returned successfully
5195        Ok(unsafe { NonEmptyVec::new_unchecked(f) })
5196    }
5197
5198    /// Compute magnitude spectrum with optional windowing, reusing cached plans.
5199    ///
5200    /// # Automatic Zero-Padding
5201    ///
5202    /// If the input signal is shorter than `n_fft`, it will be automatically
5203    /// zero-padded to the required length.
5204    ///
5205    /// # Errors
5206    /// Returns `InvalidInput` error if the input length exceeds `n_fft`.
5207    ///
5208    /// # Examples
5209    ///
5210    /// ```
5211    /// use spectrograms::*;
5212    /// use non_empty_slice::non_empty_vec;
5213    ///
5214    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
5215    /// let mut planner = FftPlanner::new();
5216    ///
5217    /// let signal = non_empty_vec![1.0; nzu!(512)];
5218    /// let magnitude = planner.magnitude_spectrum(&signal, nzu!(512), Some(WindowType::Hanning))?;
5219    ///
5220    /// assert_eq!(magnitude.len(), nzu!(257));
5221    /// # Ok(())
5222    /// # }
5223    /// ```
5224    #[inline]
5225    pub fn magnitude_spectrum(
5226        &mut self,
5227        samples: &NonEmptySlice<f64>,
5228        n_fft: NonZeroUsize,
5229        window: Option<WindowType>,
5230    ) -> SpectrogramResult<NonEmptyVec<f64>> {
5231        let power = self.power_spectrum(samples, n_fft, window)?;
5232        let power = power.iter().map(|&p| p.sqrt()).collect::<Vec<f64>>();
5233        // safety: power is non-empty since power_spectrum returned successfully
5234        Ok(unsafe { NonEmptyVec::new_unchecked(power) })
5235    }
5236}
5237
5238impl Default for FftPlanner {
5239    #[inline]
5240    fn default() -> Self {
5241        Self::new()
5242    }
5243}
5244
5245#[cfg(test)]
5246mod tests {
5247    use super::*;
5248
5249    #[test]
5250    fn test_sparse_matrix_basic() {
5251        // Create a simple 3x5 sparse matrix
5252        let mut sparse = SparseMatrix::new(3, 5);
5253
5254        // Row 0: only column 1 has value 2.0
5255        sparse.set(0, 1, 2.0);
5256
5257        // Row 1: columns 2 and 3
5258        sparse.set(1, 2, 0.5);
5259        sparse.set(1, 3, 1.5);
5260
5261        // Row 2: columns 0 and 4
5262        sparse.set(2, 0, 3.0);
5263        sparse.set(2, 4, 1.0);
5264
5265        // Test matrix-vector multiplication
5266        let input = vec![1.0, 2.0, 3.0, 4.0, 5.0];
5267        let mut output = vec![0.0; 3];
5268
5269        sparse.multiply_vec(&input, &mut output);
5270
5271        // Expected results:
5272        // Row 0: 2.0 * 2.0 = 4.0
5273        // Row 1: 0.5 * 3.0 + 1.5 * 4.0 = 1.5 + 6.0 = 7.5
5274        // Row 2: 3.0 * 1.0 + 1.0 * 5.0 = 3.0 + 5.0 = 8.0
5275        assert_eq!(output[0], 4.0);
5276        assert_eq!(output[1], 7.5);
5277        assert_eq!(output[2], 8.0);
5278    }
5279
5280    /// Build a Mel power-spectrogram plan natively at precision `T`.
5281    ///
5282    /// Mirrors `SpectrogramPlanner::mel_plan`, but is generic over the data
5283    /// scalar `T` so we can instantiate the typed pipeline at `f32` as well as
5284    /// the default `f64`. Uses the private internal builders directly.
5285    fn mel_power_plan_t<T: Sample>(
5286        params: &SpectrogramParams,
5287        mel: &MelParams,
5288    ) -> SpectrogramResult<SpectrogramPlan<Mel, Power, T>> {
5289        let stft = StftPlan::<T>::new(params)?;
5290        let mapping = FrequencyMapping::<Mel>::new_mel(params, mel)?;
5291        let scaling = AmplitudeScaling::<Power>::new(None);
5292        let freq_axis = build_frequency_axis::<Mel>(params, &mapping);
5293        let workspace = Workspace::<T>::new(stft.n_fft, stft.out_len, mapping.output_bins());
5294
5295        Ok(SpectrogramPlan {
5296            params: params.clone(),
5297            stft,
5298            mapping,
5299            scaling,
5300            freq_axis,
5301            workspace,
5302            _amp: PhantomData,
5303        })
5304    }
5305
5306    /// The high-level Mel power pipeline must compute natively in `T`: an `f32`
5307    /// run must agree with the `f64` run within a small relative tolerance,
5308    /// proving the f32 data path is real (not f64-then-downcast).
5309    #[test]
5310    fn mel_power_f32_agrees_with_f64() {
5311        let sample_rate = 16_000.0;
5312        let n = 16_000usize;
5313
5314        // Mix of two tones so several mel bands carry energy.
5315        let sig_f64: Vec<f64> = (0..n)
5316            .map(|i| {
5317                let t = i as f64 / sample_rate;
5318                (2.0 * std::f64::consts::PI * 440.0 * t).sin()
5319                    + 0.5 * (2.0 * std::f64::consts::PI * 1_500.0 * t).sin()
5320            })
5321            .collect();
5322        let sig_f32: Vec<f32> = sig_f64.iter().map(|&x| x as f32).collect();
5323
5324        let samples_f64 = NonEmptyVec::new(sig_f64).unwrap();
5325        let samples_f32 = NonEmptyVec::new(sig_f32).unwrap();
5326
5327        let stft = StftParams::new(nzu!(512), nzu!(256), WindowType::Hanning, true).unwrap();
5328        let params = SpectrogramParams::new(stft, sample_rate).unwrap();
5329        let mel = MelParams::new(nzu!(40), 0.0, 8000.0).unwrap();
5330
5331        let mut plan_f64 = mel_power_plan_t::<f64>(&params, &mel).unwrap();
5332        let mut plan_f32 = mel_power_plan_t::<f32>(&params, &mel).unwrap();
5333
5334        let spec_f64 = plan_f64.compute(&samples_f64).unwrap();
5335        let spec_f32 = plan_f32.compute(&samples_f32).unwrap();
5336
5337        assert_eq!(spec_f64.dim(), spec_f32.dim());
5338
5339        let mut max_rel = 0.0f64;
5340        let mut max_abs = 0.0f64;
5341        for (ref_val, test_val) in spec_f64.iter().zip(spec_f32.iter()) {
5342            let ref_val = *ref_val;
5343            let test_val = f64::from(*test_val);
5344            assert!(ref_val.is_finite() && test_val.is_finite());
5345            let abs = (ref_val - test_val).abs();
5346            let rel = abs / ref_val.abs().max(1e-12);
5347            max_abs = max_abs.max(abs);
5348            // Only hold the relative bound where the f64 value is large enough
5349            // to be meaningful.
5350            if ref_val.abs() > 1e-6 {
5351                max_rel = max_rel.max(rel);
5352            }
5353        }
5354
5355        // The f32 pipeline is genuinely computed in f32 end-to-end, so it carries
5356        // f32 rounding: a 512-point FFT plus power (|X|^2) and mel-band summation
5357        // accumulate to a worst-case relative deviation of ~2e-3 versus f64. A
5358        // bound well below 1% confirms the f32 path is real and correct (not a
5359        // downcast of f64 garbage, which would diverge wildly).
5360        assert!(
5361            max_rel < 5e-3,
5362            "f32 mel power deviates from f64 by relative {max_rel} (max abs {max_abs})"
5363        );
5364    }
5365
5366    #[test]
5367    fn test_sparse_matrix_zeros_ignored() {
5368        // Verify that zero values are not stored
5369        let mut sparse = SparseMatrix::new(2, 3);
5370
5371        sparse.set(0, 0, 1.0);
5372        sparse.set(0, 1, 0.0); // Should be ignored
5373        sparse.set(0, 2, 2.0);
5374
5375        // Only 2 values should be stored in row 0
5376        assert_eq!(sparse.values[0].len(), 2);
5377        assert_eq!(sparse.indices[0].len(), 2);
5378
5379        // The stored indices should be 0 and 2
5380        assert_eq!(sparse.indices[0], vec![0, 2]);
5381        assert_eq!(sparse.values[0], vec![1.0, 2.0]);
5382    }
5383
5384    #[test]
5385    fn test_loghz_matrix_sparsity() {
5386        // Verify that LogHz matrices are very sparse (1-2 non-zeros per row)
5387        let sample_rate = 16000.0;
5388        let n_fft = nzu!(512);
5389        let n_bins = nzu!(128);
5390        let f_min = 20.0;
5391        let f_max = sample_rate / 2.0;
5392
5393        let (matrix, _freqs) =
5394            build_loghz_matrix(sample_rate, n_fft, n_bins, f_min, f_max).unwrap();
5395
5396        // Each row should have at most 2 non-zero values (linear interpolation)
5397        for row_idx in 0..matrix.nrows() {
5398            let nnz = matrix.values[row_idx].len();
5399            assert!(
5400                nnz <= 2,
5401                "Row {} has {} non-zeros, expected at most 2",
5402                row_idx,
5403                nnz
5404            );
5405            assert!(nnz >= 1, "Row {} has no non-zeros", row_idx);
5406        }
5407
5408        // Total non-zeros should be close to n_bins * 2
5409        let total_nnz: usize = matrix.values.iter().map(|v| v.len()).sum();
5410        assert!(total_nnz <= n_bins.get() * 2);
5411        assert!(total_nnz >= n_bins.get()); // At least 1 per row
5412    }
5413
5414    #[test]
5415    fn test_mel_matrix_sparsity() {
5416        // Verify that Mel matrices are sparse (triangular filters)
5417        let sample_rate = 16000.0;
5418        let n_fft = nzu!(512);
5419        let n_mels = nzu!(40);
5420        let f_min = 0.0;
5421        let f_max = sample_rate / 2.0;
5422
5423        let matrix =
5424            build_mel_filterbank_matrix(sample_rate, n_fft, n_mels, f_min, f_max, MelNorm::None)
5425                .unwrap();
5426
5427        let n_fft_bins = r2c_output_size(n_fft.get());
5428
5429        // Calculate sparsity
5430        let total_nnz: usize = matrix.values.iter().map(|v| v.len()).sum();
5431        let total_elements = n_mels.get() * n_fft_bins;
5432        let sparsity = 1.0 - (total_nnz as f64 / total_elements as f64);
5433
5434        // Mel filterbanks should be >80% sparse
5435        assert!(
5436            sparsity > 0.8,
5437            "Mel matrix sparsity is only {:.1}%, expected >80%",
5438            sparsity * 100.0
5439        );
5440
5441        // Each mel filter should have significantly fewer than n_fft_bins non-zeros
5442        for row_idx in 0..matrix.nrows() {
5443            let nnz = matrix.values[row_idx].len();
5444            assert!(
5445                nnz < n_fft_bins / 2,
5446                "Mel filter {} is not sparse enough",
5447                row_idx
5448            );
5449        }
5450    }
5451}