signal_processing 0.3.0

A signal processing library.
Documentation
use core::ops::{AddAssign, MulAssign};

use num::{complex::ComplexFloat, traits::float::FloatConst, Complex, NumCast, One, Zero};
use array_math::SliceMath;
use option_trait::Maybe;

use crate::quantities::{ContainerOrSingle, List, ListOrSingle, Lists, OwnedList};

pub trait Czt<T>: Lists<T>
where
    T: ComplexFloat
{
    fn czt<R, P>(self, ratio: R, point: P) -> Self::Mapped<Complex<T::Real>>
    where
        R: Maybe<Complex<T::Real>>,
        P: Maybe<Complex<T::Real>>;
}

impl<T, L, LL> Czt<T> for L
where
    T: ComplexFloat + Into<Complex<T::Real>>,
    L: Lists<T, RowOwned: OwnedList<T, Mapped<Complex<T::Real>> = LL>, RowsMapped<LL>: Into<L::Mapped<Complex<T::Real>>>>,
    <L::RowOwned as ContainerOrSingle<T>>::Mapped<()>: List<(), Mapped<Complex<T::Real>>: Into<LL>>,
    Complex<T::Real>: AddAssign + MulAssign + MulAssign<T::Real>
{
    fn czt<R, P>(self, ratio: R, point: P) -> Self::Mapped<Complex<T::Real>>
    where
        R: Maybe<Complex<T::Real>>,
        P: Maybe<Complex<T::Real>>
    {
        let a = point.into_option()
            .map(|a| a.into())
            .unwrap_or_else(Complex::one);
        let w = ratio.into_option()
            .map(|w| w.into());

        self.map_rows_into_owned(|x| {
            let n = x.length();
            let nfft = (2*n).saturating_sub(1);
            let nfft_pow2 = nfft.next_power_of_two();

            let y_void = x.map_to_owned(|_| ());

            let w = w.unwrap_or_else(|| Complex::cis(-T::Real::TAU()/<T::Real as NumCast>::from(n).unwrap()));
            let w_sqrt = w.sqrt();
            let w2: Vec<_> = (2..=2*n).map(|i| {
                let p = i as i32 - 1 - n as i32;
                if let Some(pp) = p.checked_mul(p)
                {
                    w_sqrt.powi(pp)
                }
                else
                {
                    w_sqrt.powi(p).powi(p)
                }
            }).collect();
            let mut fw: Vec<Complex<_>> = w2[..nfft].iter()
                .map(|&w| w.recip().into())
                .collect();
            fw.resize(nfft_pow2, Complex::zero());
            fw.fft();

            let mut fg: Vec<_> = x.into_vec()
                .into_iter()
                .map(|x| x.into())
                .collect();
            let a_recip = a.recip();
            let mut apmk = Complex::one();
            for (i, g)  in fg.iter_mut()
                .enumerate()
            {
                *g *= apmk*w2[i + n - 1];
                apmk *= a_recip;
            }
            fg.resize(nfft_pow2, Complex::zero());
            fg.fft();

            let mut gg: Vec<_> = fg.into_iter()
                .zip(fw)
                .map(|(g, w)| g*w)
                .collect();
            gg.ifft();

            let mut y = gg.into_iter()
                .zip(w2)
                .skip(n.saturating_sub(1))
                .take(n)
                .map(|(g, w)| g*w);

            y_void.map_to_owned(|_| y.next().unwrap())
                .into()
        }).into()
    }
}

#[cfg(test)]
mod test
{
    use crate::transforms::fourier::Czt;

    #[test]
    fn test()
    {
        let x = [1.0, 2.0, 3.0, 4.0];

        let y = x.czt((), ());

        println!("{:?}", y);
    }
}