microspectrogram/
lib.rs

1#![no_std]
2#![allow(incomplete_features)]
3#![feature(generic_const_exprs)]
4
5use core::num::NonZeroUsize;
6
7use microfft::Complex32;
8
9pub use crate::window::Window;
10use crate::{
11    utils::{detrend_constant, overlapping_windows, Assert, IsTrue},
12    window::get_window,
13};
14
15mod fft;
16mod utils;
17mod window;
18
19pub struct Spectrogram<const N: usize> {
20    fs: usize,
21    win: Window,
22}
23
24impl<const N: usize> Spectrogram<N>
25where
26    Assert<{ (N == 4) | (N == 8) | (N == 16) }>: IsTrue,
27{
28    pub fn new(fs: usize, win: Window) -> Self {
29        Self { fs, win }
30    }
31
32    pub fn compute<const X: usize>(&self, data: &[f32; X]) -> [[f32; X / (N / 2) - 1]; 1 + N / 2]
33    where
34        [(); N + 1]: Sized,
35    {
36        let overlap = N / 2;
37        let wins = overlapping_windows(
38            data,
39            NonZeroUsize::new(N).unwrap(),
40            NonZeroUsize::new(overlap).unwrap(),
41        );
42
43        let mut seg = [0.0; N];
44        let mut col = [0.0; 1 + N / 2];
45        let mut out = [[0.0; X / (N / 2) - 1]; 1 + N / 2];
46
47        for (i, data) in wins.enumerate() {
48            seg.copy_from_slice(data);
49            process_segment::<N>(&mut seg, &mut col, self.fs, self.win);
50            col.into_iter().enumerate().for_each(|(j, v)| {
51                out[j][i] = v;
52            });
53        }
54        out
55    }
56}
57
58fn process_segment<const N: usize>(
59    arr: &mut [f32; N],
60    result: &mut [f32; 1 + N / 2],
61    fs: usize,
62    win: Window,
63) where
64    [(); N + 1]: Sized,
65{
66    let win = get_window::<N>(win, false);
67    detrend_constant(arr);
68    arr.iter_mut().zip(win).for_each(|(x, w)| {
69        *x *= w;
70    });
71
72    // This is the output of the _fft_helper funtion in scipy
73    let mut result_c = [Complex32::new(0.0, 0.0); 1 + N / 2];
74
75    // SAFETY: We know that `arr` has the right length (compatible with `result`)
76    // because of the arguments of `process_segment`.
77    unsafe { fft::fft(arr, &mut result_c) };
78
79    let mut win_sq_sum = 0.0;
80    win.iter().for_each(|x| {
81        win_sq_sum += x * x;
82    });
83
84    // scaling == 'density'
85    let scale = 1.0 / (fs as f32 * win_sq_sum);
86
87    result_c.iter_mut().for_each(|x| {
88        let conj = Complex32::new(x.re, -x.im);
89        *x *= conj * scale * 2.0;
90    });
91
92    result_c[0] /= 2.0;
93    result_c[result_c.len() - 1] /= 2.0;
94
95    result.iter_mut().zip(result_c).for_each(|(r, rc)| {
96        *r = rc.re;
97    });
98}
99
100#[cfg(test)]
101mod test {
102    use core::ops;
103
104    use num_traits::ToPrimitive;
105
106    use super::*;
107
108    fn all_close<T, const N: usize>(a: &[T; N], b: &[T; N]) -> bool
109    where
110        T: ops::Sub<Output = T> + ToPrimitive + Copy,
111    {
112        const EPSILON: f32 = 1e-5;
113        for i in 0..N {
114            if (a[i] - b[i]).to_f32().unwrap() > EPSILON {
115                return false;
116            }
117        }
118        true
119    }
120
121    fn all_close2d<T, const N: usize, const M: usize>(a: &[[T; N]; M], b: &[[T; N]; M]) -> bool
122    where
123        T: ops::Sub<Output = T> + ToPrimitive + Copy,
124    {
125        const EPSILON: f32 = 1e-5;
126        for i in 0..M {
127            for j in 0..N {
128                if (a[i][j] - b[i][j]).to_f32().unwrap() > EPSILON {
129                    return false;
130                }
131            }
132        }
133        true
134    }
135
136    #[test]
137    #[allow(clippy::excessive_precision)]
138    fn test_segment() {
139        let mut input = [
140            0.65600173, 0.63951888, 0.62348176, 0.60798841, 0.59066761, 0.57568833, 0.56063477,
141            0.54556932, 0.53043256, 0.51324547, 0.4984534, 0.48411299, 0.46861073, 0.45568773,
142            0.4428301, 0.42829952,
143        ];
144        let expected_scipy = [
145            8.22254870e-04,
146            1.87558778e-02,
147            9.88275690e-04,
148            4.09807662e-05,
149            1.52053553e-05,
150            8.80124399e-07,
151            1.94954969e-06,
152            4.08357001e-07,
153            6.61331734e-10,
154        ];
155        let mut result = Default::default();
156        process_segment(&mut input, &mut result, 1, Window::Hann);
157        assert!(all_close(&result, &expected_scipy));
158    }
159
160    #[test]
161    #[allow(clippy::excessive_precision)]
162    fn test_spectrogram() {
163        let x = [
164            0.65600173, 0.63951888, 0.62348176, 0.60798841, 0.59066761, 0.57568833, 0.56063477,
165            0.54556932, 0.53043256, 0.51324547, 0.4984534, 0.48411299, 0.46861073, 0.45568773,
166            0.4428301, 0.42829952, 0.41565583, 0.40287546, 0.38830922, 0.37554074, 0.36229681,
167            0.34613488, 0.32802069, 0.30968066,
168        ];
169        let expected = [
170            [8.22254870e-04, 4.46185932e-04],
171            [1.87558778e-02, 1.44202233e-02],
172            [9.88275690e-04, 6.00894689e-04],
173            [4.09807662e-05, 3.58241931e-05],
174            [1.52053553e-05, 5.18071551e-06],
175            [8.80124399e-07, 8.84743592e-06],
176            [1.94954969e-06, 1.82243933e-07],
177            [4.08357001e-07, 2.45994606e-07],
178            [6.61331734e-10, 8.14788254e-09],
179        ];
180
181        let spec = Spectrogram::<16>::new(1, Window::Hann);
182        let out = spec.compute(&x);
183        assert!(all_close2d(&out, &expected));
184    }
185}
microspectrogram/lib.rs

microspectrogram/
lib.rs
