csrk 1.1.4

/*! Wendland Kernel structs and methods

See Eq. 4.21 from 
<a href=https://direct.mit.edu/books/oa-monograph/2320/Gaussian-Processes-for-Machine-Learning target="_blank">
Rasmussen and Williams (2006)</a>: 
*/

// ---- Imports ----
// Third Party
use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use sprs::{TriMat};
// Local
use crate::spatial_hash::SpatialHash;

/// Wendland Kernels should have methods which return information
pub trait Wendland {
    /// The order of mean square differentiability
    fn msq_order(&self) -> i32;
    /// The dimensionality of the input data
    fn read_ndim(&self) -> i32;
    /// Staticmethod: The "j" integer from Eq. 4.21)
    fn calc_msq_diff_int(ndim: i32, order: i32) -> i32 {
        // Note that this is integer division and the desired behavior
        // is flooring ndim
        (ndim / 2) + order + 1
    }
    /// Staticmethod: j + order
    fn calc_msq_diff_power(order: i32, msq_diff_int: i32) -> i32 {
        order + msq_diff_int}
    /// Evaluate the kernel basis function
    fn eval_sgl(&self, r: f64, r2: f64) -> f64;
}
/// Order 0 Basis for Wendland Kernel
pub struct WendlandBasisMSQ0{
    ndim: i32,
    #[allow(unused)]
    msq_diff_int: i32,
    msq_diff_power: i32,
}
// General implementations for WendlandBasisMSQ0
impl WendlandBasisMSQ0 {
    pub fn new(ndim: usize, order: i32) -> Self {
        let msq_diff_int = Self::calc_msq_diff_int(ndim as i32, order);
        Self {
            ndim: ndim as i32,
            msq_diff_int,
            msq_diff_power: Self::calc_msq_diff_power(order, msq_diff_int),
        }
    }
}
// Wendland trait implementation for WendlandBasisMSQ0
impl Wendland for WendlandBasisMSQ0 {
    fn msq_order(&self) -> i32 {0}
    fn read_ndim(&self) -> i32 {self.ndim.clone()}
    fn eval_sgl(&self, r: f64, _r2: f64) -> f64 {
        // check usage
        if r >= 1.0 {return 0.0;}
        // Compute kernel normally
        (1. - r).max(0.0).powi(self.msq_diff_power)
    }
}

/// Order 1 Basis for Wendland Kernel
pub struct WendlandBasisMSQ1{
    ndim: i32,
    #[allow(unused)]
    msq_diff_int: i32,
    msq_diff_power: i32,
}
// General implementations for WendlandBasisMSQ1
impl WendlandBasisMSQ1 {
    pub fn new(ndim: usize, order: i32) -> Self {
        let msq_diff_int = Self::calc_msq_diff_int(ndim as i32, order);
        Self {
            ndim: ndim as i32,
            msq_diff_int,
            msq_diff_power: Self::calc_msq_diff_power(order, msq_diff_int)
        }
    }
}
// Wendland trait implementation for WendlandBasisMSQ1
impl Wendland for WendlandBasisMSQ1 {
    fn msq_order(&self) -> i32 {1}
    fn read_ndim(&self) -> i32 {self.ndim.clone()}
    fn eval_sgl(&self, r: f64, _r2: f64) -> f64 {
        // check usage
        if r >= 1.0 {return 0.0;}
        // Compute kernel normally
        (1. - r).max(0.0).powi(self.msq_diff_power) * 
            ((f64::from(self.msq_diff_power) * r) + 1.0)
    }
}


/// Order 2 Basis for Wendland Kernel
pub struct WendlandBasisMSQ2{
    ndim: i32,
    #[allow(unused)]
    msq_diff_int: i32,
    msq_diff_power: i32,
    j_poly_r1: f64,
    j_poly_r2: f64,
}
// General implementations for WendlandBasisMSQ2
impl WendlandBasisMSQ2 {
    pub fn new(ndim: usize, order: i32) -> Self {
        let msq_diff_int = Self::calc_msq_diff_int(ndim as i32, order);
        Self {
            ndim: ndim as i32,
            msq_diff_int,
            msq_diff_power: Self::calc_msq_diff_power(order, msq_diff_int),
            j_poly_r1: Self::calc_j_poly_r1(msq_diff_int) as f64,
            j_poly_r2: Self::calc_j_poly_r2(msq_diff_int) as f64,
        }
    }
    /// Staticmethod for constructing the j_poly_r1 coefficient
    fn calc_j_poly_r1(msq_diff_int: i32) -> f64 {
        f64::from((3 * msq_diff_int) + 6)
    }
    /// Staticmethod for constructing the j_poly_r2 coefficient
    fn calc_j_poly_r2(msq_diff_int: i32) -> f64 {
        f64::from(msq_diff_int.pow(2) + 4*msq_diff_int + 3)
    }
}
// Wendland trait implementation for WendlandBasisMSQ2
impl Wendland for WendlandBasisMSQ2 {
    fn msq_order(&self) -> i32 {2}
    fn read_ndim(&self) -> i32 {self.ndim.clone()}
    fn eval_sgl(&self, r: f64, r2: f64) -> f64 {
        // check usage
        if r >= 1.0 {return 0.0;}
        // Compute kernel normally
        (1. - r).max(0.0).powi(self.msq_diff_power) * (
            (self.j_poly_r2 * r2) + 
            (self.j_poly_r1 * r) +
            3.0
        ) / 3.0
    }
}

/// Order 3 Basis for Wendland Kernel
pub struct WendlandBasisMSQ3{
    ndim: i32,
    #[allow(unused)]
    msq_diff_int: i32,
    msq_diff_power: i32,
    j_poly_r1: f64,
    j_poly_r2: f64,
    j_poly_r3: f64,
}
// General implementations for WendlandBasisMSQ2
impl WendlandBasisMSQ3 {
    pub fn new(ndim: usize, order: i32) -> Self {
        let msq_diff_int = Self::calc_msq_diff_int(ndim as i32, order);
        Self {
            ndim: ndim as i32,
            msq_diff_int,
            msq_diff_power: Self::calc_msq_diff_power(order, msq_diff_int),
            j_poly_r1: Self::calc_j_poly_r1(msq_diff_int) as f64,
            j_poly_r2: Self::calc_j_poly_r2(msq_diff_int) as f64,
            j_poly_r3: Self::calc_j_poly_r3(msq_diff_int) as f64,
        }
    }
    /// Staticmethod for constructing the j_poly_r1 coefficient
    fn calc_j_poly_r1(msq_diff_int: i32) -> f64 {
        f64::from((15 * msq_diff_int) + 45)
    }
    /// Staticmethod for constructing the j_poly_r2 coefficient
    fn calc_j_poly_r2(msq_diff_int: i32) -> f64 {
        f64::from(6 * msq_diff_int.pow(2) + 36*msq_diff_int + 45)
    }
    /// Staticmethod for constructing the j_poly_r3 coefficient
    fn calc_j_poly_r3(msq_diff_int: i32) -> f64 {
        f64::from(
            msq_diff_int.pow(3) + 
            9 * msq_diff_int.pow(2) + 
            23*msq_diff_int + 15
            )
    }
}
// Wendland trait implementation for WendlandBasisMSQ3
impl Wendland for WendlandBasisMSQ3 {
    fn msq_order(&self) -> i32 {3}
    fn read_ndim(&self) -> i32 {self.ndim.clone()}
    fn eval_sgl(&self, r: f64, r2: f64) -> f64 {
        // check usage
        if r >= 1.0 {return 0.0;}
        // Compute kernel normally
        let r3 = r2 * r;
        (1. - r).max(0.0).powi(self.msq_diff_power)  * (
            (self.j_poly_r3 * r3) + 
            (self.j_poly_r2 * r2) + 
            (self.j_poly_r1 * r) + 
            15.
        ) / 15.
    }
}

/// Enum for the Wendland Kernel structs
pub enum WendlandKernel {
    Q0(WendlandBasisMSQ0),
    Q1(WendlandBasisMSQ1),
    Q2(WendlandBasisMSQ2),
    Q3(WendlandBasisMSQ3),
}
/// General constructor for the Wendland Kernel
impl WendlandKernel {
    pub fn new(ndim: usize, msq_order: i32) -> Self {
        match msq_order {
            0 => WendlandKernel::Q0(WendlandBasisMSQ0::new(ndim, msq_order)),
            1 => WendlandKernel::Q1(WendlandBasisMSQ1::new(ndim, msq_order)),
            2 => WendlandKernel::Q2(WendlandBasisMSQ2::new(ndim, msq_order)),
            3 => WendlandKernel::Q3(WendlandBasisMSQ3::new(ndim, msq_order)),
            _ => unreachable!(),
        }
    }
}
// Wendland implementation for the WendlandKernel enum
impl Wendland for WendlandKernel {
    fn msq_order(&self)     -> i32 {
        match self {
            WendlandKernel::Q0(k) => k.msq_order(),
            WendlandKernel::Q1(k) => k.msq_order(),
            WendlandKernel::Q2(k) => k.msq_order(),
            WendlandKernel::Q3(k) => k.msq_order(),
        }
    }
    fn read_ndim(&self)     -> i32 {
        match self {
            WendlandKernel::Q0(k) => k.read_ndim(),
            WendlandKernel::Q1(k) => k.read_ndim(),
            WendlandKernel::Q2(k) => k.read_ndim(),
            WendlandKernel::Q3(k) => k.read_ndim(),
        }
    }
    fn eval_sgl(&self, r:f64, r2:f64) -> f64 {
        match self {
            WendlandKernel::Q0(k) => k.eval_sgl(r, r2),
            WendlandKernel::Q1(k) => k.eval_sgl(r, r2),
            WendlandKernel::Q2(k) => k.eval_sgl(r, r2),
            WendlandKernel::Q3(k) => k.eval_sgl(r, r2),
        }
    }
}

/// General kernel struct
pub struct Kernel {
    basis: WendlandKernel,
    scale: Array1<f64>,
    whitenoise: f64
}
impl Kernel {
    /// Return a new kernel object
    pub fn new(
        ndim        : usize,
        msq_order   : i32,
        scale_view  : ArrayView1<f64>,
        whitenoise  : f64,
    ) -> Self {
        // Get a basis for the WendlandKernel
        let basis = WendlandKernel::new(ndim, msq_order);
        let scale = scale_view.as_standard_layout().into_owned();
        Self {
            basis,
            scale,
            whitenoise,
        }
    }
    /// Call the basis eval single
    pub fn eval_sgl(
        &self, r:f64, r2:f64
    ) -> f64 {
        self.basis.eval_sgl(r, r2)
    }
    /// Return a reference to the order
    pub fn msq_order(
        &self,
    ) -> i32 {
        self.basis.msq_order()
    }
    /// Return a reference to the whitenoise
    pub fn whitenoise(
        &self,
    ) -> f64 {
        self.whitenoise
    }
    /// Return a reference to the scale
    pub fn scale(
        &self,
    ) -> Array1<f64> {
        self.scale.clone()
    }
    /// Scale unscaled training or evaluation data (sgl)
    pub fn scale_sgl(&self, sample_point: ArrayView1<f64>) -> Array1<f64> {
        &sample_point / &self.scale
    }
    /// Scale unscaled training or evaluation data (arr)
    pub fn scale_arr(&self, sample_points: ArrayView2<f64>) -> Array2<f64> {
        let mut scaled_points = sample_points.as_standard_layout().into_owned();
        scaled_points /= &self.scale;
        scaled_points
    }
    /**
    Construct a sparse evaluated kernel by looping i and j
    */
    pub fn naive_kernel_construction(
        &self,
        scaled_training_points: ArrayView2<f64>,
        training_error: ArrayView1<f64>,
    ) -> TriMat<f64> {

        // Identify dimensions
        let nsample = scaled_training_points.nrows();
        // Initialize training_kernel
        let mut training_kernel = TriMat::<f64>::new((nsample,nsample));

        // Assemble the sparse kernel
        for i in 0..nsample {
            // Fill in diagonals
            // Wendland of r==0 always returns 1.0
            let diag_cell_value = training_error[i].powi(2) + self.whitenoise + 1e-10;
            training_kernel.add_triplet(i,i,1.0 + diag_cell_value);

            // Get row for point 1
            let xi = scaled_training_points.row(i);

            // Fill in off-diagonals
            for j in 0..i {
                // Identify x2
                let xj = scaled_training_points.row(j);
                // Calculate r2 using iterator
                let r2: f64 = xi.iter()
                    .zip(xj.iter())
                    .map(|(ki,kj)| (ki - kj).powi(2))
                    .sum();

                // Check for skip condition
                // iff sqrt(r2) > 1 then r2 > 1
                if r2 >= 1.0 {continue;}
                // sqrt r
                let r = r2.sqrt();
                // Evaluate kernel
                let val: f64 = self.eval_sgl(r, r2);
                training_kernel.add_triplet(i,j,val);
                training_kernel.add_triplet(j,i,val);
            }
        }
        training_kernel
    }
    /**
    Construct a sparse evaluated kernel by visiting close points
    */
    pub fn hashed_kernel_construction(
        &self,
        scaled_training_points: ArrayView2<f64>,
        training_error: ArrayView1<f64>,
        training_hash: &SpatialHash,
    ) -> TriMat<f64> {
        // Identify dimensions
        let nsample = scaled_training_points.nrows();
        // Initialize training_kernel
        let mut training_kernel = TriMat::<f64>::new((nsample,nsample));
        // Assemble the sparse kernel
        for i in 0..nsample {
            // Fill in diagonals
            // Wendland of r==0 always returns 1.0
            let diag_cell_value = training_error[i].powi(2) + self.whitenoise + 1e-10;
            training_kernel.add_triplet(i,i,1.0 + diag_cell_value);
            
            // Identify the training point in question
            let xi = scaled_training_points.row(i);

            // Visit neighbors
            training_hash.for_each_neighbor(xi.as_slice().unwrap(), |j| {
                // upper triangle only
                if j <= i {return;}
                // Identify xj value
                let xj = scaled_training_points.row(j);
                
                // Calculate r2
                let r2: f64 =
                    xi.iter().zip(
                    xj.iter())
                    .map(|(ki,kj)| (ki-kj).powi(2))
                    .sum();
                // Check for skip condition
                // iff sqrt(r2) > 1 then r2 > 1
                if r2 >= 1.0 {return;}
                // sqrt r
                let r = r2.sqrt();
                let val: f64 = self.eval_sgl(r, r2);
                training_kernel.add_triplet(i,j,val);
                training_kernel.add_triplet(j,i,val);
            });
        }
        training_kernel
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Instant;
    use crate::spatial_hash::SpatialHash;
    use ndarray::{arr1, Array2};

    #[test]
    fn timing_kernel_build_naive_vs_hashed() {
        // ---- build a 2D grid of training points ----
        let nx = 150;
        let ny = 150;
        let h = 0.2; // spacing < 1.0 so neighbors fall inside support
        let ndim = 2;
        let scale = arr1(&vec![1.;ndim]);

        let mut x_train_vec: Vec<Vec<f64>> = Vec::new();
        for i in 0..nx {
            for j in 0..ny {
                x_train_vec.push(vec![i as f64 * h, j as f64 * h]);
            }
        }
        // Flatten for array construction
        let x_train_flat: Vec<f64> = x_train_vec.iter()
            .flat_map(|row| row.iter().copied())
            .collect();

        let nsample = x_train_vec.len();

        println!("points: {}", nsample);
        let x_train = Array2::from_shape_vec((nsample, ndim), x_train_flat).unwrap();

        // dummy data (kernel build only depends on geometry + errors)
        let y_err = arr1(&vec![0.0; nsample]);
        let whitenoise = 1e-10;

        // kernel parameters (pick one representative case)
        let msq_order = 2;
        let mykernel = Kernel::new(
            ndim,
            msq_order,
            scale.view(),
            whitenoise
        );
        // Create hash for training data
        let hash = SpatialHash::build(&x_train_vec, 1.0);

        // ---- naive build ----
        let t0 = Instant::now();
        let kn = mykernel.naive_kernel_construction(
            x_train.view(),
            y_err.view(),
        );
        let t_naive = t0.elapsed();

        // ---- hashed build ----
        let t1 = Instant::now();
        let kh = mykernel.hashed_kernel_construction(
            x_train.view(),
            y_err.view(),
            &hash,
        );
        let t_hash = t1.elapsed();

        let nnz_naive = kn.nnz();
        let nnz_hash = kh.nnz();

        println!("naive build : {:?}", t_naive);
        println!("hashed build: {:?}", t_hash);
        println!("speedup     : {:.2}x", t_naive.as_secs_f64() / t_hash.as_secs_f64());
        println!("nnz naive   : {}", nnz_naive);
        println!("nnz hash    : {}", nnz_hash);

        // ---- correctness sanity check ----
        assert_eq!(
            nnz_naive, nnz_hash,
            "hashed kernel missed or added interactions"
        );
    }
}