lasprs 0.14.0

use super::fft::FFT;
use crate::*;
use ndarray::prelude::*;
use num::pow::Pow;
use realfft::{RealFftPlanner, RealToComplex};
use reinterpret::reinterpret_slice;
use std::mem::MaybeUninit;
use std::sync::Arc;

const SMALL_NUMBER: Flt = 1e-80;

/// Cross power spectra, which is a 3D array, with the following properties:
///
/// - The first index is the frequency index, starting at DC, ending at nfft/2.
/// - The second, and third index result in `[i,j]` = C_ij = p_i  * conj(p_j)
///
#[cfg_attr(feature = "python-bindings", gen_stub_pyclass, pyclass(from_py_object))]
#[derive(Debug, Clone)]
pub struct CPSResult(Array3<Cflt>);

impl CPSResult {
    /// Creates a new `CPSResult` from underlying `Array3<Cflt>`.
    pub fn new(result: Array3<Cflt>) -> Self {
        CPSResult(result)
    }
    /// Returns the underlying `Array3<Cflt>`.
    pub fn into_inner(self) -> Array3<Cflt> {
        self.0
    }
    /// Scale the power spectrum by the given sensitivity value for each channel.
    ///
    /// # Arguments
    ///
    /// * `sens` - Sensitivity values for each channel.
    ///
    /// # Panics
    ///
    /// Panics if the number of sensitivity values does not match the number of
    /// channels.
    pub fn apply_sensitivities(&mut self, sens: &[Flt]) {
        assert_eq!(sens.len(), self.0.shape()[1]);
        debug_assert!(self.0.shape()[2] == self.0.shape()[1]);
        for (i, mut coli) in &mut self.0.axis_iter_mut(Axis(1)).enumerate() {
            for (j, mut colj) in coli.axis_iter_mut(Axis(1)).enumerate() {
                colj.map_inplace(|x| *x = *x / sens[j] / sens[i]);
            }
        }
    }

    /// Returns a reference to the underlying `Array3<Cflt>`.
    pub fn inner(&self) -> &Array3<Cflt> {
        &self.0
    }

    /// Returns a mutable reference to the underlying `Array3<Cflt>`.
    pub fn inner_mut(&mut self) -> &mut Array3<Cflt> {
        &mut self.0
    }
    /// Whether there not  is any data (The array has a non-zero shape).
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
    /// The number of channels in the result.
    pub fn nchannels(&self) -> usize {
        self.0.shape()[1]
    }

    /// The number of FFT bins of the result.
    pub fn nfft(&self) -> usize {
        self.0.shape()[0]
    }

    /// Returns the autopower for a single channel, as a array of real values
    /// (imaginary part is zero and is stripped off).
    ///
    /// # Args
    ///
    /// - `ch` - The channel number to compute autopower for.
    pub fn ap(&self, ch: usize) -> Array1<Flt> {
        // Slice out one value for all frequencies, map to only real part, and
        // return.
        self.0.slice(s![.., ch, ch]).mapv(|f| f.re)
    }

    /// Returns the autopower for a single channel in dB, relative to the
    /// provided reference value. The reference valuse should be provided on a
    /// linear scale. Adds a tiny offset of 0f32.next_up() to avoid log(0) for
    /// situations where the autopower is zero.
    ///
    /// - `ch` - The channel number to compute autopower for.
    /// - `ref_value` - Reference value for the dB value.
    pub fn ap_dB(&self, ch: usize, ref_value: StrictlyPositive) -> Array1<Flt> {
        let ref_value_sq = *ref_value * *ref_value;
        let mut ap = self.ap(ch);
        // Convert in place to dB relative to the reference value.
        ap.mapv_inplace(|p| 10. * Flt::log10(Flt::max(p, SMALL_NUMBER) / ref_value_sq));
        ap
    }

    /// Returns the coherence between two channels as a function of frequency. The
    /// coherence is defined as the squared magnitude of the cross-spectral
    /// density divided by the product of the individual spectral densities.
    ///
    /// # Args
    ///
    /// - `ch1` - The channel number of the first channel.
    /// - `ch2` - The channel number of the second channel.
    ///
    /// Returns: gamma_ij: Magnitude squared coherence between channels `ch1`
    /// and `ch2`.
    pub fn coherence(&self, ch1: usize, ch2: usize) -> Array1<Flt> {
        // Is there a way to avoid the allocation of `res`?
        let mut res = Array1::zeros(self.0.shape()[0]);
        azip!((r in &mut res,
            c00 in self.0.slice(s![..,ch1,ch1]),
            c11 in self.0.slice(s![..,ch2,ch2]),
            c01 in self.0.slice(s![..,ch1,ch2]))
            {
                let den  = Flt::max(c00.re*c11.re, SMALL_NUMBER);
                let num = (c01*c01.conj()).re;
                *r =num/den;
            }
        );
        res
    }

    /// Returns the transfer function estimationfrom `chi` to `chj`, that is ~
    /// Pj/Pi,  as a array of complex numbers.
    ///
    /// # Args
    ///
    /// - `chi` - The channel number of the *denominator*
    /// - `chj` - The channel number of the *numerator*
    /// - `chRef` - Optional, a reference channel that has the lowest noise. If
    ///   not given, the average of the two autopowers is used, which gives
    ///   always a worse result than when two a low noise reference channel is
    ///   used.
    ///
    pub fn tf(&self, chi: usize, chj: usize, chRef: Option<usize>) -> Array1<Cflt> {
        match chRef {
            None => {
                let cij = self.0.slice(s![.., chi, chj]);
                let cii = self.0.slice(s![.., chi, chi]);
                let cjj = self.0.slice(s![.., chj, chj]);
                Zip::from(cij)
                    .and(cii)
                    .and(cjj)
                    .par_map_collect(|cij, cii, cjj| 0.5 * (cij.conj() / cii + cjj / cij))
            }
            Some(chr) => {
                let cir = self.0.slice(s![.., chi, chr]);
                let cjr = self.0.slice(s![.., chj, chr]);

                Zip::from(cir)
                    .and(cjr)
                    .par_map_collect(|cir, cjr| cjr / cir)
            }
        }
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl CPSResult {
    /// Returns the autopower for a single channel, as a Python array of real values.
    ///
    /// # Args
    ///
    /// - `ch` - The channel number to compute autopower for. Panics if out of
    /// bounds.
    #[pyo3(name = "ap")]
    fn ap_py<'py>(&self, py: Python<'py>, ch: usize) -> Bound<'py, PyArray1<Flt>> {
        self.ap(ch).to_pyarray(py)
    }

    /// Returns the transfer function from `chi` to `chj`, that is ~ Pj/Pi for
    /// an acoustic frequency response, as a array of complex
    /// numbers.
    ///
    /// # Args
    ///
    /// - `chi` - The channel number of the *denominator*
    /// - `chj` - The channel number of the *numerator*
    /// - `chRef` - Optional, a reference channel that has the lowest noise. If
    ///   not given, the average of the two autopowers is used, which gives
    ///   always a worse result than when two a low noise reference channel is
    ///   used.
    ///
    #[pyo3(name = "tf")]
    fn tf_py<'py>(
        &self,
        py: Python<'py>,
        chi: usize,
        chj: usize,
        chRef: Option<usize>,
    ) -> Bound<'py, PyArray1<Cflt>> {
        self.tf(chi, chj, chRef).to_pyarray(py)
    }
    /// Direct access to the cross-power spectrum components, Cij = Pi * conj(Pj)
    ///
    /// # Args
    ///
    /// - `chi` - The channel number for Pi
    /// - `chj` - The channel number for Pj
    #[pyo3(name = "C")]
    fn C_py<'py>(&self, py: Python<'py>, chi: usize, chj: usize) -> Bound<'py, PyArray1<Cflt>> {
        self.inner().slice(s![.., chi, chj]).to_pyarray(py)
    }

    /// Return the full cross-power spectrum as a 3D array, Cij = Pi * conj(Pj)
    #[pyo3(name = "Cfull")]
    fn Cfull_py<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray3<Cflt>> {
        self.inner().to_pyarray(py)
    }

    /// See `CPSResult::ap_dB()`
    #[pyo3(name = "ap_dB")]
    fn ap_dB_py<'py>(
        &self,
        py: Python<'py>,
        ch: usize,
        ref_value: StrictlyPositive,
    ) -> Bound<'py, PyArray1<Flt>> {
        self.ap_dB(ch, ref_value).into_pyarray(py)
    }

    /// See `CPSResult::coherence()`
    #[pyo3(name = "coherence")]
    fn coherence_py<'py>(
        &self,
        py: Python<'py>,
        ch1: usize,
        ch2: usize,
    ) -> Bound<'py, PyArray1<Flt>> {
        self.coherence(ch1, ch2).into_pyarray(py)
    }
}

/// Single-sided (cross)power spectra estimator, that uses a Windowed FFT to
/// estimate cross-power spectra. Window functions are documented in the
/// `window` module. Note that directly using this power spectra estimator is
/// generally not useful as it is basically the periodogram estimator, with its
/// high variance.
///
/// This power spectrum estimator is instead used as a building block for for
/// example the computations of spectrograms, or Welch' method of spectral
/// estimation.
///
#[derive(Debug, Clone)]
pub struct PowerSpectra {
    /// Window used in estimator. The actual Window in here is normalized with
    /// the square root of the Window power. This safes one division when
    /// processing time data.
    pub window_normalized: Window,

    ffts: Vec<FFT>,

    // Time-data buffer used for multiplying signals with Window
    timedata: Array2<Flt>,
    // Frequency domain buffer used for storage of signal FFt's in inbetween stage
    freqdata: Array2<Cflt>,

    // Result storage
    res: CPSResult,
}

impl PowerSpectra {
    /// Returns the FFT length used in power spectra computations
    #[inline]
    pub fn nfft(&self) -> usize {
        self.window_normalized.win.len()
    }
    /// Create new power spectra estimator. Uses FFT size from window length
    ///
    /// # Panics
    ///
    /// - If win.len() != nfft
    /// - if nfft == 0
    ///
    /// # Args
    ///
    /// - `window` - A `Window` struct, from which NFFT is also used.
    ///
    pub fn newFromWindow(mut window: Window) -> PowerSpectra {
        let nfft = window.win.len();
        let win_pwr = window.win.mapv(|w| w.powi(2)).sum() / (nfft as Flt);
        let sqrt_win_pwr = Flt::sqrt(win_pwr);
        window.win.mapv_inplace(|v| v / sqrt_win_pwr);

        assert!(nfft > 0);
        assert!(nfft.is_multiple_of(2));

        let mut planner = RealFftPlanner::<Flt>::new();
        let fft = planner.plan_fft_forward(nfft);

        let Fft = FFT::new(fft);
        let res = CPSResult::new(Array3::zeros((0, 0, 0)));

        PowerSpectra {
            window_normalized: window,
            ffts: vec![Fft],
            timedata: Array2::zeros((nfft, 1)),
            freqdata: Array2::zeros((nfft / 2 + 1, 1)),
            res,
        }
    }

    /// Compute FFTs of input channel data. Stores the scaled FFT data in
    /// self.freqdata.
    fn compute_ffts(&mut self, timedata: ArrayView2<Flt>) {
        assert!(timedata.nrows() > 0);
        let (n, nch) = timedata.dim();
        let nfft = self.nfft();
        assert!(n == nfft);

        // Make sure enough fft engines are available
        while nch > self.ffts.len() {
            self.ffts
                .push(self.ffts.last().expect("FFT's should not be empty").clone());
            self.freqdata
                .push_column(Ccol::from_vec(vec![Cflt::new(0., 0.); nfft / 2 + 1]).view())
                .unwrap();
            self.timedata.push_column(Dcol::zeros(nfft).view()).unwrap();
        }

        assert!(n == self.nfft());
        assert!(n == self.window_normalized.win.len());

        // Multiply signals with window function, and compute fft's for each channel
        Zip::from(timedata.axis_iter(Axis(1)))
            .and(self.timedata.axis_iter_mut(Axis(1)))
            .and(&mut self.ffts)
            .and(self.freqdata.axis_iter_mut(Axis(1)))
            .par_for_each(|time_in, mut time_tmp_storage, fft, mut freq| {
                let DC = time_in.mean().unwrap();

                azip!((t in &mut time_tmp_storage, &tin in time_in, &win in &self.window_normalized.win) {
                // Substract DC value from time data, as this leaks into
                // positive frequencies due to windowing.
                // Multiply with window and copy over to local time data buffer
                    *t=(tin-DC)*win});

                fft.process(&time_tmp_storage, &mut freq);
                freq[0] = DC + 0. * I;
            });
    }

    /// Compute cross power spectra from input time data. First axis is
    /// frequency, second axis is channel i, third axis is channel j.
    ///
    /// # Panics
    ///
    /// - When `timedata.nrows() != self.nfft()`
    ///
    /// # Args
    ///
    /// * `tdata` - Input time data. This is a 2D array, where the first axis is
    ///   time and the second axis is the channel number.
    ///
    ///  # Returns
    ///
    ///  - 3D complex array of signal cross-powers with the following shape
    ///    (nfft/2+1,timedata.ncols(), timedata.ncols()). Its content is:
    ///    [freq_index, chi, chj] = crosspower: chi*conj(chj)
    ///
    pub fn compute<'a, 'b, T>(&'b mut self, tdata: T) -> &'b mut CPSResult
    where
        T: AsArray<'a, Flt, Ix2>,
    {
        let tdata = tdata.into();
        let nfft = self.nfft();
        let clen = nfft / 2 + 1;
        if tdata.nrows() != nfft {
            panic!("Invalid timedata length! Should be equal to nfft={nfft}");
        }
        let nchannels = tdata.ncols();

        // Compute fft of input data, and store in self.freqdata
        self.compute_ffts(tdata);
        let fd = &self.freqdata;
        let fdconj = fd.mapv(|c| c.conj());

        let result = self.res.inner_mut();
        {
            // Check if result array needs to be resized
            let required_shape = (clen, nchannels, nchannels).f();
            let required_dim = required_shape.raw_dim();
            if result.raw_dim() != *required_dim {
                *result = Array3::zeros(required_shape);
            }
        }

        // Loop over result axis one and channel i IN PARALLEL
        Zip::from(result.axis_iter_mut(Axis(1)))
            .and(fd.axis_iter(Axis(1)))
            .par_for_each(|mut out, chi| {
                // out: channel i of output 3D array, channel j all
                // chi: channel i
                Zip::from(out.axis_iter_mut(Axis(1)))
                    .and(fdconj.axis_iter(Axis(1)))
                    .for_each(|mut out, chj| {
                        // out: channel i, j
                        // chj: channel j conjugated
                        Zip::from(&mut out)
                            .and(chi)
                            .and(chj)
                            .for_each(|out, chi, chjc| {
                                // Loop over frequency components
                                *out = 0.5 * chi * chjc;
                            });

                        // The DC component has no 0.5 correction, as it only
                        // occurs ones in a (double-sided) power spectrum. So
                        // here we undo the 0.5 of 4 lines above here.
                        out[0] *= 2.;
                        out[clen - 1] *= 2.;
                    });
            });
        &mut self.res
    }
}

#[cfg(test)]
mod test {
    use approx::{abs_diff_eq, assert_relative_eq, assert_ulps_eq, ulps_eq};
    // For absolute value
    use num::complex::ComplexFloat;

    /// Generate a sine wave at the order i
    fn generate_sinewave(nfft: usize, order: usize) -> Dcol {
        Dcol::from_iter(
            (0..nfft).map(|i| Flt::sin(i as Flt / (nfft) as Flt * order as Flt * 2. * pi)),
        )
    }
    /// Generate a sine wave at the order i
    fn generate_cosinewave(nfft: usize, order: usize) -> Dcol {
        Dcol::from_iter(
            (0..nfft).map(|i| Flt::cos(i as Flt / (nfft) as Flt * order as Flt * 2. * pi)),
        )
    }

    use crate::math::randNormal;

    use super::*;
    #[test]
    /// Test whether DC part of single-sided FFT has right properties
    fn test_fft_DC() {
        const nfft: usize = 10;
        let rect = Window::new(WindowType::Rect, nfft);
        let mut ps = PowerSpectra::newFromWindow(rect);

        let td = Dmat::ones((nfft, 1));

        ps.compute_ffts(td.view());
        let fd = &ps.freqdata;
        // println!("{:?}", fd);
        assert_relative_eq!(fd[(0, 0)].re, 1.);
        assert_relative_eq!(fd[(0, 0)].im, 0.);
        let abs_fneq0 = fd.slice(s![1.., 0]).sum();
        assert_relative_eq!(abs_fneq0.re, 0.);
        assert_relative_eq!(abs_fneq0.im, 0.);
    }

    /// Test whether AC part of single-sided FFT has right properties
    #[test]
    fn test_fft_AC() {
        const nfft: usize = 256;
        let rect = Window::new(WindowType::Rect, nfft);
        let mut ps = PowerSpectra::newFromWindow(rect);

        // Start with a time signal
        let mut t: Dmat = Dmat::default((nfft, 0));
        t.push_column(generate_sinewave(nfft, 1).view()).unwrap();
        // println!("{:?}", t);

        ps.compute_ffts(t.view());
        let fd = &ps.freqdata;
        // println!("{:?}", fd);
        assert_relative_eq!(fd[(0, 0)].re, 0., epsilon = Flt::EPSILON * nfft as Flt);
        assert_relative_eq!(fd[(0, 0)].im, 0., epsilon = Flt::EPSILON * nfft as Flt);

        assert_relative_eq!(fd[(1, 0)].re, 0., epsilon = Flt::EPSILON * nfft as Flt);
        assert_ulps_eq!(fd[(1, 0)].im, -1., epsilon = Flt::EPSILON * nfft as Flt);

        // Sum of all terms at frequency index 2 to ...
        let sum_higher_freqs_abs = Cflt::abs(fd.slice(s![2.., 0]).sum());
        assert_ulps_eq!(
            sum_higher_freqs_abs,
            0.,
            epsilon = Flt::EPSILON * nfft as Flt
        );
    }

    /// Thest whether power spectra scale properly. Signals with amplitude of 1
    /// should come back with a power of 0.5. DC offsets should come in as
    /// value^2 at frequency index 0.
    #[test]
    fn test_ps_scale() {
        const nfft: usize = 124;
        let rect = Window::new(WindowType::Rect, nfft);
        let mut ps = PowerSpectra::newFromWindow(rect);

        // Start with a time signal
        let mut t: Dmat = Dmat::default((nfft, 0));
        t.push_column(generate_cosinewave(nfft, 1).view()).unwrap();
        let dc_component = 0.25;
        let dc_power = dc_component.pow(2);
        t.mapv_inplace(|t| t + dc_component);

        let power = ps.compute(t.view());
        let power = power.inner();
        assert_relative_eq!(
            power[(0, 0, 0)].re,
            dc_power,
            epsilon = Flt::EPSILON * nfft as Flt
        );
        assert_relative_eq!(
            power[(1, 0, 0)].re,
            0.5,
            epsilon = Flt::EPSILON * nfft as Flt
        );
        assert_relative_eq!(
            power[(1, 0, 0)].im,
            0.0,
            epsilon = Flt::EPSILON * nfft as Flt
        );
    }

    // Test parseval's theorem for some random data
    #[test]
    fn test_parseval() {
        const nfft: usize = 512;
        let rect = Window::new(WindowType::Rect, nfft);
        let mut ps = PowerSpectra::newFromWindow(rect);

        // Start with a time signal
        let t: Dmat = randNormal((nfft, 1));

        let tavg = t.sum() / (nfft as Flt);
        let t_dc_power = tavg.powi(2);
        // println!("dc power in time domain: {:?}", t_dc_power);

        let signal_pwr = t.mapv(|t| t.powi(2)).sum() / (nfft as Flt);
        // println!("Total signal power in time domain: {:?} ", signal_pwr);

        let power = ps.compute(t.view());
        let power = power.inner();
        // println!("freq domain power: {:?}", power);

        let fpower = power.sum().abs();

        assert_ulps_eq!(
            t_dc_power,
            power[(0, 0, 0)].abs(),
            epsilon = Flt::EPSILON * (nfft as Flt).powi(2)
        );
        assert_ulps_eq!(
            signal_pwr,
            fpower,
            epsilon = Flt::EPSILON * (nfft as Flt).powi(2)
        );
    }

    // Test parseval's theorem for some random data
    #[test]
    fn test_parseval_with_window() {
        // A sufficiently high value is required here, to show that it works.
        const nfft: usize = 2usize.pow(20);
        let window = Window::new(WindowType::Hann, nfft);
        // let window = Window::new(WindowType::Rect, nfft);
        let mut ps = PowerSpectra::newFromWindow(window);

        // Start with a time signal
        let t: Dmat = randNormal((nfft, 1));

        let tavg = t.sum() / (nfft as Flt);
        let t_dc_power = tavg.powi(2);
        // println!("dc power in time domain: {:?}", t_dc_power);

        let signal_pwr = t.mapv(|t| t.powi(2)).sum() / (nfft as Flt);
        // println!("Total signal power in time domain: {:?} ", signal_pwr);

        let power = ps.compute(t.view());
        let power = power.inner();
        // println!("freq domain power: {:?}", power);

        let fpower = power.sum().abs();

        assert_ulps_eq!(
            t_dc_power,
            power[(0, 0, 0)].abs(),
            epsilon = Flt::EPSILON * (nfft as Flt).powi(2)
        );

        // This one fails when nfft is too short.
        assert_ulps_eq!(signal_pwr, fpower, epsilon = 2e-2);
    }
}