audio_lab/signal/ops/
spectrogram.rs

1use ndarray::{Array1, Array3, Zip};
2use num::complex::Complex64;
3
4use crate::signal::Spectrogram;
5
6fn one_sided_scale(spec: &Spectrogram, bin: usize) -> f64 {
7    let last_bin = spec.num_freq_bins().saturating_sub(1);
8    if bin == 0 || (spec.window_size().is_multiple_of(2) && bin == last_bin) {
9        1.0
10    } else {
11        2.0
12    }
13}
14
15/// Apply `value` to every bin and multiply by the per-bin one-sided scale
16/// divided by `denom`. The scales are computed once per bin, not per element.
17fn scaled_per_bin(
18    spec: &Spectrogram,
19    denom: f64,
20    value: impl Fn(&Complex64) -> f64,
21) -> Array3<f64> {
22    let scales = Array1::from_shape_fn(spec.num_freq_bins(), |bin| {
23        one_sided_scale(spec, bin) / denom
24    });
25    Zip::from(spec.data())
26        .and_broadcast(&scales)
27        .map_collect(|c, &scale| value(c) * scale)
28}
29
30fn to_db_in_place(mut values: Array3<f64>, factor: f64, floor_linear: f64) -> Array3<f64> {
31    values.mapv_inplace(|v| factor * v.max(floor_linear).log10());
32    values
33}
34
35impl Spectrogram {
36    pub fn raw_magnitude(&self) -> Array3<f64> {
37        self.data().map(|c| c.norm())
38    }
39
40    pub fn raw_magnitude_db(&self, floor_db: f64) -> Array3<f64> {
41        let floor_linear = 10f64.powf(floor_db / 20.0);
42        self.data()
43            .map(|c| 20.0 * c.norm().max(floor_linear).log10())
44    }
45
46    pub fn amplitude_spectrum(&self) -> Array3<f64> {
47        let coherent_sum = self.normalization().coherent_gain() * self.window_size() as f64;
48        scaled_per_bin(self, coherent_sum, |c| c.norm())
49    }
50
51    pub fn amplitude_spectrum_db(&self, floor_db: f64) -> Array3<f64> {
52        let floor_linear = 10f64.powf(floor_db / 20.0);
53        to_db_in_place(self.amplitude_spectrum(), 20.0, floor_linear)
54    }
55
56    pub fn power_spectrum(&self) -> Array3<f64> {
57        let denom = self.window_size() as f64 * self.normalization().window_energy();
58        scaled_per_bin(self, denom, |c| c.norm_sqr())
59    }
60
61    pub fn power_spectrum_db(&self, floor_db: f64) -> Array3<f64> {
62        let floor_linear = 10f64.powf(floor_db / 10.0);
63        to_db_in_place(self.power_spectrum(), 10.0, floor_linear)
64    }
65
66    pub fn power_spectral_density(&self) -> Array3<f64> {
67        let denom = self.sample_rate() * self.normalization().window_energy();
68        scaled_per_bin(self, denom, |c| c.norm_sqr())
69    }
70
71    pub fn power_spectral_density_db(&self, floor_db: f64) -> Array3<f64> {
72        let floor_linear = 10f64.powf(floor_db / 10.0);
73        to_db_in_place(self.power_spectral_density(), 10.0, floor_linear)
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use approx::assert_abs_diff_eq;
80    use ndarray::Array3;
81    use num::complex::Complex64;
82
83    use crate::dsp::{WindowFn, generate_periodic_window};
84    use crate::signal::Spectrogram;
85
86    fn make_spectrogram(window_size: usize) -> Spectrogram {
87        let num_freq_bins = window_size / 2 + 1;
88        let data = Array3::from_shape_fn((2, 3, num_freq_bins), |(ch, frame, bin)| {
89            Complex64::new(
90                (ch + 1) as f64 * (bin as f64 + 0.5),
91                frame as f64 - bin as f64 * 0.25,
92            )
93        });
94        Spectrogram::new(data, 8_000.0, window_size, window_size / 2, WindowFn::Hann).unwrap()
95    }
96
97    #[test]
98    fn scaled_spectra_match_per_element_formulas() {
99        // Even window: DC and Nyquist bins are not doubled. Odd window: only DC.
100        for window_size in [8, 9] {
101            let spec = make_spectrogram(window_size);
102            let last_bin = spec.num_freq_bins() - 1;
103            let scale = |bin: usize| {
104                if bin == 0 || (window_size % 2 == 0 && bin == last_bin) {
105                    1.0
106                } else {
107                    2.0
108                }
109            };
110            // Gains of the Hann window the spectrogram was built with.
111            let window = generate_periodic_window(WindowFn::Hann, window_size);
112            let window_energy: f64 = window.iter().map(|w| w * w).sum();
113            let coherent_sum = window.sum();
114            let power_denom = window_size as f64 * window_energy;
115            let psd_denom = 8_000.0 * window_energy;
116
117            let amplitude = spec.amplitude_spectrum();
118            let amplitude_db = spec.amplitude_spectrum_db(-200.0);
119            let power = spec.power_spectrum();
120            let power_db = spec.power_spectrum_db(-200.0);
121            let psd = spec.power_spectral_density();
122            let psd_db = spec.power_spectral_density_db(-200.0);
123
124            for ((ch, frame, bin), c) in spec.data().indexed_iter() {
125                let idx = [ch, frame, bin];
126                let expected_amplitude = c.norm() * scale(bin) / coherent_sum;
127                let expected_power = c.norm_sqr() * scale(bin) / power_denom;
128                let expected_psd = c.norm_sqr() * scale(bin) / psd_denom;
129
130                assert_abs_diff_eq!(amplitude[idx], expected_amplitude, epsilon = 1e-12);
131                assert_abs_diff_eq!(power[idx], expected_power, epsilon = 1e-12);
132                assert_abs_diff_eq!(psd[idx], expected_psd, epsilon = 1e-12);
133                assert_abs_diff_eq!(
134                    amplitude_db[idx],
135                    20.0 * expected_amplitude.log10(),
136                    epsilon = 1e-9
137                );
138                assert_abs_diff_eq!(power_db[idx], 10.0 * expected_power.log10(), epsilon = 1e-9);
139                assert_abs_diff_eq!(psd_db[idx], 10.0 * expected_psd.log10(), epsilon = 1e-9);
140            }
141        }
142    }
143}
audio_lab/signal/ops/spectrogram.rs

audio_lab/signal/ops/
spectrogram.rs
