funspace 0.5.0

//! # Orthogonal Chebyshev Base
use crate::enums::{BaseKind, TransformKind};
use crate::traits::{
    BaseElements, BaseFromOrtho, BaseGradient, BaseMatOpDiffmat, BaseMatOpLaplacian,
    BaseMatOpStencil, BaseSize, BaseTransform,
};
use crate::types::{FloatNum, ScalarNum};
use ndarray::{s, Array2};
use rustdct::{Dct1, DctPlanner};
use std::f64::consts::PI;
use std::sync::Arc;

/// # Container for chebyshev base
#[derive(Clone)]
pub struct Chebyshev<A> {
    /// Number of coefficients in physical space
    n: usize,
    /// Number of coefficients in spectral space ( equal to *n* for `Chebyshev` )
    m: usize,
    /// `DCT`-Plan
    plan_dct: Arc<dyn Dct1<A>>,
}

impl<A: FloatNum> Chebyshev<A> {
    /// Creates a new Basis.
    ///
    /// # Arguments
    /// * `n` - Length of array's dimension which shall live in chebyshev space.
    ///
    /// # Panics
    /// Panics when input type cannot be cast from f64.
    ///
    /// # Examples
    /// ```
    /// use funspace::chebyshev::Chebyshev;
    /// let cheby = Chebyshev::<f64>::new(10);
    /// ```
    #[must_use]
    pub fn new(n: usize) -> Self {
        let mut planner = DctPlanner::<A>::new();
        let dct1 = planner.plan_dct1(n);
        Self {
            n,
            m: n,
            plan_dct: Arc::clone(&dct1),
        }
    }

    // /// Chebyshev points of the second kind. $[-1, 1]$
    // /// $$$
    // /// x = - cos( pi*k/(npts - 1) )
    // /// $$$
    // #[allow(clippy::cast_precision_loss)]
    // fn chebyshev_nodes_2nd_kind(n: usize) -> Vec<A> {
    //     (0..n)
    //         .map(|k| -A::one() * A::from_f64(PI * k as f64 / (n - 1) as f64).unwrap().cos())
    //         .collect::<Vec<A>>()
    // }

    /// Chebyshev nodes of the second kind on intervall $[-1, 1]$
    /// $$$
    /// x = - cos( pi*k/(npts - 1) )
    /// $$$
    #[allow(clippy::cast_precision_loss)]
    fn chebyshev_nodes_2nd_kind(n: usize) -> Vec<A> {
        let m = (n - 1) as f64;
        (0..n)
            .map(|k| {
                let arg = A::from_f64(PI * (m - 2. * k as f64) / (2. * m)).unwrap();
                -A::one() * arg.sin()
            })
            .collect::<Vec<A>>()
    }

    /// Returns grid points
    #[must_use]
    pub fn nodes(n: usize) -> Vec<A> {
        Self::chebyshev_nodes_2nd_kind(n)
    }

    /// Differentation Matrix see [`chebyshev::dmsuite::diffmat_chebyshev`
    ///
    /// # Panics
    /// - Num conversion fails
    /// - deriv > 4: Not implemented
    #[must_use]
    #[allow(clippy::cast_precision_loss)]
    pub fn _dmat(n: usize, deriv: usize) -> Array2<A> {
        let mut dmat = Array2::<f64>::zeros((n, n));
        if deriv == 1 {
            for p in 0..n {
                for q in p + 1..n {
                    if (p + q) % 2 != 0 {
                        dmat[[p, q]] = (q * 2) as f64;
                    }
                }
            }
        } else if deriv == 2 {
            for p in 0..n {
                for q in p + 2..n {
                    if (p + q) % 2 == 0 {
                        dmat[[p, q]] = (q * (q * q - p * p)) as f64;
                    }
                }
            }
        } else if deriv == 3 {
            for p in 0..n {
                let p2 = p * p;
                for q in p + 3..n {
                    let q2 = q * q;
                    if (p + q) % 2 != 0 {
                        dmat[[p, q]] =
                            (q * (q2 * (q2 - 2) - 2 * q2 * p2 + p2 * p2 - 2 * p2 + 1)) as f64 / 4.;
                    }
                }
            }
        } else if deriv == 4 {
            for p in 0..n {
                let (p2, p4) = (p * p, p * p * p * p);
                for q in p + 4..n {
                    let (q2, q4) = (q * q, q * q * q * q);
                    if (p + q) % 2 == 0 {
                        dmat[[p, q]] = (q
                            * (q2 * (q2 - 4) * (q2 - 4) + 3 * q2 * p4 - p2 * p4 + 8 * p4
                                - 16 * p2
                                - 3 * q4 * p2)) as f64
                            / 24.;
                    }
                }
            }
        } else {
            todo!()
        }
        for d in dmat.slice_mut(s![0, ..]).iter_mut() {
            *d *= 0.5;
        }
        dmat.mapv(|elem| A::from_f64(elem).unwrap())
    }
    /// Pseudoinverse matrix of chebyshev spectral
    /// differentiation matrices
    ///
    /// When preconditioned with the pseudoinverse Matrix,
    /// systems become banded and thus efficient to solve.
    ///
    /// Literature:
    /// Sahuck Oh - An Efficient Spectral Method to Solve Multi-Dimensional
    /// Linear Partial Different Equations Using Chebyshev Polynomials
    ///
    /// Output:
    /// ndarray (n x n) matrix, acts in spectral space
    ///
    /// # Panics
    /// - Num conversion fails
    /// - deriv > 4: Not implemented
    #[must_use]
    #[allow(clippy::cast_precision_loss)]
    pub fn _pinv(n: usize, deriv: usize) -> Array2<A> {
        // assert!(
        //     !(deriv > 2),
        //     "pinv does only support deriv's 1 & 2, got {}",
        //     deriv
        // );
        let mut pinv = Array2::<f64>::zeros([n, n]);
        if deriv == 1 {
            pinv[[1, 0]] = 1.;
            for i in 2..n {
                pinv[[i, i - 1]] = 1. / (2. * i as f64); // diag - 1
            }
            for i in 1..n - 2 {
                pinv[[i, i + 1]] = -1. / (2. * i as f64); // diag + 1
            }
        } else if deriv == 2 {
            pinv[[2, 0]] = 0.25;
            for i in 3..n {
                pinv[[i, i - 2]] = 1. / (4 * i * (i - 1)) as f64; // diag - 2
            }
            for i in 2..n - 2 {
                pinv[[i, i]] = -1. / (2 * (i * i - 1)) as f64; // diag 0
            }
            for i in 2..n - 4 {
                pinv[[i, i + 2]] = 1. / (4 * i * (i + 1)) as f64; // diag + 2
            }
        } else if deriv == 3 {
            // diag - 3
            for i in 3..n {
                let d = 8 * i * (i - 1) * (i - 2);
                pinv[[i, i - 3]] = 1. / d as f64;
            }
            pinv[[3, 0]] *= 2.;
            // diag - 1
            for i in 3..n - 2 {
                let d = 8 * (i + 1) * (i - 2) * i;
                pinv[[i, i - 1]] = -3. / d as f64;
            }
            // diag + 1
            for i in 3..n - 4 {
                let d = 8 * (i - 1) * (i + 2) * i;
                pinv[[i, i + 1]] = 3. / d as f64;
            }
            // diag + 3
            for i in 3..n - 6 {
                let d = 8 * i * (i + 1) * (i + 2);
                pinv[[i, i + 3]] = -1. / d as f64;
            }
        } else if deriv == 4 {
            // diag - 4
            for i in 4..n {
                let d = 16 * i * (i - 1) * (i - 2) * (i - 3);
                pinv[[i, i - 4]] = 1. / d as f64;
            }
            pinv[[4, 0]] *= 2.;
            // diag - 2
            for i in 4..n - 2 {
                let d = 4 * (i - 3) * (i - 1) * i * (i + 1);
                pinv[[i, i - 2]] = -1. / d as f64;
            }
            // diag + 0
            for i in 4..n - 4 {
                let d = 8 * (i - 2) * (i - 1) * (i + 1) * (i + 2);
                pinv[[i, i]] = 3. / d as f64;
            }
            // diag + 2
            for i in 4..n - 6 {
                let d = 4 * (i - 1) * i * (i + 1) * (i + 3);
                pinv[[i, i + 2]] = -1. / d as f64;
            }
            // diag + 4
            for i in 4..n - 8 {
                let d = 16 * i * (i + 1) * (i + 2) * (i + 3);
                pinv[[i, i + 4]] = 1. / d as f64;
            }
        } else {
            todo!()
        }
        //pinv
        pinv.mapv(|elem| A::from_f64(elem).unwrap())
    }

    /// Returns eye matrix, where the n ( = deriv) upper rows
    ///
    /// # Panics
    /// Num conversion fails
    #[must_use]
    pub fn _pinv_eye(n: usize, deriv: usize) -> Array2<A> {
        let pinv_eye = Array2::<f64>::eye(n).slice(s![deriv.., ..]).to_owned();
        pinv_eye.mapv(|elem| A::from_f64(elem).unwrap())
    }
}

impl<A: FloatNum> BaseSize for Chebyshev<A> {
    /// Size in physical space
    fn len_phys(&self) -> usize {
        self.n
    }

    /// Size in spectral space
    fn len_spec(&self) -> usize {
        self.m
    }

    /// Size of orthogonal space
    fn len_orth(&self) -> usize {
        self.m
    }
}

impl<A: FloatNum> BaseElements for Chebyshev<A> {
    /// Real valued scalar type
    type RealNum = A;

    /// Return kind of base
    fn base_kind(&self) -> BaseKind {
        BaseKind::Chebyshev
    }

    /// Return kind of transform
    fn transform_kind(&self) -> TransformKind {
        TransformKind::R2r
    }

    /// Coordinates in physical space
    fn coords(&self) -> Vec<A> {
        Chebyshev::nodes(self.len_phys())
    }
}

impl<A: FloatNum> BaseMatOpDiffmat for Chebyshev<A> {
    /// Scalar type of matrix
    type NumType = A;

    /// Explicit differential operator $ D $
    ///
    /// Matrix-based version of [`BaseGradient::gradient()`]
    fn diffmat(&self, deriv: usize) -> Array2<Self::NumType> {
        assert!(deriv > 0);
        Self::_dmat(self.n, deriv)
    }

    /// Explicit inverse of differential operator $ D^* $
    ///
    /// Returns ``(D_pinv, I_pinv)``, where `D_pinv` is the pseudoinverse
    /// and ``I_pinv`` the corresponding pseudoidentity matrix, such
    /// that
    ///
    /// ```text
    /// D_pinv @ D = I_pinv
    /// ```
    ///
    /// Can be used as a preconditioner.
    fn diffmat_pinv(&self, deriv: usize) -> (Array2<Self::NumType>, Array2<Self::NumType>) {
        assert!(deriv > 0);
        (Self::_pinv(self.n, deriv), Self::_pinv_eye(self.n, deriv))
    }
}

impl<A: FloatNum> BaseMatOpStencil for Chebyshev<A> {
    /// Scalar type of matrix
    type NumType = A;

    /// Transformation stencil composite -> orthogonal space
    fn stencil(&self) -> Array2<Self::NumType> {
        Array2::<A>::eye(self.len_spec())
    }

    /// Inverse of transformation stencil
    fn stencil_inv(&self) -> Array2<Self::NumType> {
        Array2::<A>::eye(self.len_spec())
    }
}

impl<A: FloatNum> BaseMatOpLaplacian for Chebyshev<A> {
    /// Scalar type of matrix
    type NumType = A;

    /// Laplacian $ L $
    fn laplacian(&self) -> Array2<Self::NumType> {
        self.diffmat(2)
    }

    /// Pseudoinverse matrix of Laplacian $ L^{-1} $
    ///
    /// Returns pseudoinverse and pseudoidentity,i.e
    /// ``(D_pinv, I_pinv)``
    ///
    /// ```text
    /// D_pinv @ D = I_pinv
    /// ``
    fn laplacian_pinv(&self) -> (Array2<Self::NumType>, Array2<Self::NumType>) {
        self.diffmat_pinv(2)
    }
}

impl<A, T> BaseFromOrtho<T> for Chebyshev<A>
where
    A: FloatNum,
    T: ScalarNum,
{
    /// Composite space coefficients -> Orthogonal space coefficients
    fn to_ortho_slice(&self, indata: &[T], outdata: &mut [T]) {
        for (y, x) in outdata.iter_mut().zip(indata.iter()) {
            *y = *x;
        }
    }

    /// Orthogonal space coefficients -> Composite space coefficients
    fn from_ortho_slice(&self, indata: &[T], outdata: &mut [T]) {
        for (y, x) in outdata.iter_mut().zip(indata.iter()) {
            *y = *x;
        }
    }
}

impl<A, T> BaseGradient<T> for Chebyshev<A>
where
    A: FloatNum,
    T: ScalarNum,
{
    /// Differentiate Vector `n_times` using recurrence relation
    /// of chebyshev polynomials.
    ///
    /// # Panics
    /// Float conversion fails (unlikely)
    ///
    /// # Example
    /// Differentiate Chebyshev
    /// ```
    /// use funspace::traits::BaseGradient;
    /// use funspace::chebyshev::Chebyshev;
    /// use funspace::utils::approx_eq;
    /// let mut ch = Chebyshev::<f64>::new(4);
    /// let indata: Vec<f64> = vec![1., 2., 3., 4.];
    /// let mut outdata: Vec<f64> = vec![0.; 4];
    /// // Differentiate twice
    /// ch.gradient_slice(&indata, &mut outdata, 2);
    /// approx_eq(&outdata, &vec![12., 96.,  0.,  0.]);
    /// ```
    fn gradient_slice(&self, indata: &[T], outdata: &mut [T], n_times: usize) {
        assert!(outdata.len() == self.m);
        // Copy over
        for (y, x) in outdata.iter_mut().zip(indata.iter()) {
            *y = *x;
        }
        // Differentiate
        let two = T::one() + T::one();
        for _ in 0..n_times {
            // Forward
            unsafe {
                *outdata.get_unchecked_mut(0) = *outdata.get_unchecked(1);
                for i in 1..self.m - 1 {
                    *outdata.get_unchecked_mut(i) =
                        two * T::from_usize(i + 1).unwrap() * *outdata.get_unchecked(i + 1);
                }
                *outdata.get_unchecked_mut(self.m - 1) = T::zero();
                // Reverse
                for i in (1..self.m - 2).rev() {
                    *outdata.get_unchecked_mut(i) =
                        *outdata.get_unchecked(i) + *outdata.get_unchecked(i + 2);
                }
                *outdata.get_unchecked_mut(0) =
                    *outdata.get_unchecked(0) + *outdata.get_unchecked(2) / two;
            }
        }
    }
}

impl<A: FloatNum> BaseTransform for Chebyshev<A> {
    type Physical = A;

    type Spectral = A;

    fn forward_slice(&self, indata: &[Self::Physical], outdata: &mut [Self::Spectral]) {
        // Check input
        assert!(indata.len() == self.len_phys());
        assert!(indata.len() == outdata.len());

        // Copy and correct input data
        let cor = (A::one() + A::one()) * A::one() / A::from(self.n - 1).unwrap();

        // Reverse indata since it is defined on $[-1, 1]$, instead of $[1, -1]$
        for (y, x) in outdata.iter_mut().zip(indata.iter().rev()) {
            *y = *x * cor;
        }
        // Transform via dct
        self.plan_dct.process_dct1(outdata);

        // Correct first and last coefficient
        let half = A::from_f64(0.5).unwrap();
        outdata[0] *= half;
        outdata[self.n - 1] *= half;
    }

    fn backward_slice(&self, indata: &[Self::Spectral], outdata: &mut [Self::Physical]) {
        // Check input
        assert!(indata.len() == self.len_spec());
        assert!(indata.len() == outdata.len());

        // Copy and correct input data
        // Multiplying with [-1, 1, -1, ...] reverses order of output
        // data, such that outdata is defined on $[-1, 1]$ intervall
        let two = A::one() + A::one();
        for (i, (y, x)) in outdata.iter_mut().zip(indata.iter()).enumerate() {
            if i % 2 == 0 {
                *y = *x;
            } else {
                *y = *x * -A::one();
            }
        }
        outdata[0] *= two;
        outdata[self.m - 1] *= two;

        // Transform via dct
        self.plan_dct.process_dct1(outdata);
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::utils::approx_eq;

    #[test]
    fn test_chebyshev_transform_1() {
        let cheby = Chebyshev::<f64>::new(4);
        let indata = vec![1., 2., 3., 4.];
        let mut outdata = vec![0.; 4];
        cheby.forward_slice(&indata, &mut outdata);
        approx_eq(&outdata, &vec![2.5, 1.33333333, 0., 0.16666667]);
    }

    #[test]
    fn test_chebyshev_pinv() {
        let n = 27;
        let ch = Chebyshev::<f64>::new(n);
        // construct b4 by matmul
        let (b2, _) = ch.diffmat_pinv(2);
        let mut b4_v2 = b2.dot(&b2);
        b4_v2.slice_mut(ndarray::s![..4, ..]).fill(0.);
        b4_v2.slice_mut(ndarray::s![.., n - 4..]).fill(0.);
        // get b4 directly
        let (b4, _) = ch.diffmat_pinv(4);
        // compare
        for (x, y) in b4_v2.iter().zip(b4.iter()) {
            assert!((x - y).abs() < 1e-6);
        }
    }
}