microspectrogram 0.1.0

#![no_std]
#![allow(incomplete_features)]
#![feature(generic_const_exprs)]

use core::num::NonZeroUsize;

use microfft::Complex32;

pub use crate::window::Window;
use crate::{
    utils::{detrend_constant, overlapping_windows, Assert, IsTrue},
    window::get_window,
};

mod fft;
mod utils;
mod window;

pub struct Spectrogram<const N: usize> {
    fs: usize,
    win: Window,
}

impl<const N: usize> Spectrogram<N>
where
    Assert<{ (N == 4) | (N == 8) | (N == 16) }>: IsTrue,
{
    pub fn new(fs: usize, win: Window) -> Self {
        Self { fs, win }
    }

    pub fn compute<const X: usize>(&self, data: &[f32; X]) -> [[f32; X / (N / 2) - 1]; 1 + N / 2]
    where
        [(); N + 1]: Sized,
    {
        let overlap = N / 2;
        let wins = overlapping_windows(
            data,
            NonZeroUsize::new(N).unwrap(),
            NonZeroUsize::new(overlap).unwrap(),
        );

        let mut seg = [0.0; N];
        let mut col = [0.0; 1 + N / 2];
        let mut out = [[0.0; X / (N / 2) - 1]; 1 + N / 2];

        for (i, data) in wins.enumerate() {
            seg.copy_from_slice(data);
            process_segment::<N>(&mut seg, &mut col, self.fs, self.win);
            col.into_iter().enumerate().for_each(|(j, v)| {
                out[j][i] = v;
            });
        }
        out
    }
}

fn process_segment<const N: usize>(
    arr: &mut [f32; N],
    result: &mut [f32; 1 + N / 2],
    fs: usize,
    win: Window,
) where
    [(); N + 1]: Sized,
{
    let win = get_window::<N>(win, false);
    detrend_constant(arr);
    arr.iter_mut().zip(win).for_each(|(x, w)| {
        *x *= w;
    });

    // This is the output of the _fft_helper funtion in scipy
    let mut result_c = [Complex32::new(0.0, 0.0); 1 + N / 2];

    // SAFETY: We know that `arr` has the right length (compatible with `result`)
    // because of the arguments of `process_segment`.
    unsafe { fft::fft(arr, &mut result_c) };

    let mut win_sq_sum = 0.0;
    win.iter().for_each(|x| {
        win_sq_sum += x * x;
    });

    // scaling == 'density'
    let scale = 1.0 / (fs as f32 * win_sq_sum);

    result_c.iter_mut().for_each(|x| {
        let conj = Complex32::new(x.re, -x.im);
        *x *= conj * scale * 2.0;
    });

    result_c[0] /= 2.0;
    result_c[result_c.len() - 1] /= 2.0;

    result.iter_mut().zip(result_c).for_each(|(r, rc)| {
        *r = rc.re;
    });
}

#[cfg(test)]
mod test {
    use core::ops;

    use num_traits::ToPrimitive;

    use super::*;

    fn all_close<T, const N: usize>(a: &[T; N], b: &[T; N]) -> bool
    where
        T: ops::Sub<Output = T> + ToPrimitive + Copy,
    {
        const EPSILON: f32 = 1e-5;
        for i in 0..N {
            if (a[i] - b[i]).to_f32().unwrap() > EPSILON {
                return false;
            }
        }
        true
    }

    fn all_close2d<T, const N: usize, const M: usize>(a: &[[T; N]; M], b: &[[T; N]; M]) -> bool
    where
        T: ops::Sub<Output = T> + ToPrimitive + Copy,
    {
        const EPSILON: f32 = 1e-5;
        for i in 0..M {
            for j in 0..N {
                if (a[i][j] - b[i][j]).to_f32().unwrap() > EPSILON {
                    return false;
                }
            }
        }
        true
    }

    #[test]
    #[allow(clippy::excessive_precision)]
    fn test_segment() {
        let mut input = [
            0.65600173, 0.63951888, 0.62348176, 0.60798841, 0.59066761, 0.57568833, 0.56063477,
            0.54556932, 0.53043256, 0.51324547, 0.4984534, 0.48411299, 0.46861073, 0.45568773,
            0.4428301, 0.42829952,
        ];
        let expected_scipy = [
            8.22254870e-04,
            1.87558778e-02,
            9.88275690e-04,
            4.09807662e-05,
            1.52053553e-05,
            8.80124399e-07,
            1.94954969e-06,
            4.08357001e-07,
            6.61331734e-10,
        ];
        let mut result = Default::default();
        process_segment(&mut input, &mut result, 1, Window::Hann);
        assert!(all_close(&result, &expected_scipy));
    }

    #[test]
    #[allow(clippy::excessive_precision)]
    fn test_spectrogram() {
        let x = [
            0.65600173, 0.63951888, 0.62348176, 0.60798841, 0.59066761, 0.57568833, 0.56063477,
            0.54556932, 0.53043256, 0.51324547, 0.4984534, 0.48411299, 0.46861073, 0.45568773,
            0.4428301, 0.42829952, 0.41565583, 0.40287546, 0.38830922, 0.37554074, 0.36229681,
            0.34613488, 0.32802069, 0.30968066,
        ];
        let expected = [
            [8.22254870e-04, 4.46185932e-04],
            [1.87558778e-02, 1.44202233e-02],
            [9.88275690e-04, 6.00894689e-04],
            [4.09807662e-05, 3.58241931e-05],
            [1.52053553e-05, 5.18071551e-06],
            [8.80124399e-07, 8.84743592e-06],
            [1.94954969e-06, 1.82243933e-07],
            [4.08357001e-07, 2.45994606e-07],
            [6.61331734e-10, 8.14788254e-09],
        ];

        let spec = Spectrogram::<16>::new(1, Window::Hann);
        let out = spec.compute(&x);
        assert!(all_close2d(&out, &expected));
    }
}