csrk 1.1.4

/*! Compactly Supported Radial Kernel -- Gaussian Process crate

This work was deeply inspired by Section 4.2 of
<a href=https://direct.mit.edu/books/oa-monograph/2320/Gaussian-Processes-for-Machine-Learning target="_blank">
Rasmussen and Williams (2006)</a>: 
*/
// Module imports
use sprs::CsMat;
use sprs_ldl::{LdlNumeric};
use hdf5_metno::{File, Result as H5Result};
use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis};

// Local imports
// random number generation
pub mod prng;
pub use crate::prng::PCG64Stream;
// Kernel
pub mod kernel;
pub use crate::kernel::{Wendland};
pub use crate::kernel::{WendlandKernel};
pub use crate::kernel::{Kernel};
// spatial hash mod
pub mod spatial_hash;
pub use crate::spatial_hash::SpatialHash;
// realizations
pub mod realization;
pub use crate::realization::GPRealization;

// GP struct
/**
Gaussian Process struct
*/
pub struct GP {
    pub ndim: usize,
    pub kernel: Kernel,
    training_points: Array2<f64>,   // shape: (nsample, ndim)
    training_values: Array1<f64>,   // shape: (nsample,)
    training_error: Array1<f64>,    // shape: (nsample,)
    training_kernel: CsMat<f64>,
    training_hash: SpatialHash,
    alpha: Array1<f64>,             // shape: (nsample,)
    #[allow(non_snake_case)]
    lnL: f64,
    ldl_numeric: Option<LdlNumeric<f64, usize>>,
}

impl GP {
    pub fn train(
        training_points: ArrayView2<f64>,
        training_values: ArrayView1<f64>,
        training_error: ArrayView1<f64>,
        scale: ArrayView1<f64>,
        whitenoise: f64,
        msq_order: i32,
    ) -> GP {
        /*!
        Create a new Gaussian Process struct and train it
        */
        // Identify ndim
        let ndim = training_points.ncols();
        // Initialize kernel 
        let kernel = Kernel::new(ndim, msq_order, scale, whitenoise);
        // Initialize scaled input
        let scaled_training_points: Array2<f64> = 
            kernel.scale_arr(training_points);
        // Convert inputs to C-contiguous owned arrays
        let training_values = training_values.as_standard_layout().into_owned();
        let training_error = training_error.as_standard_layout().into_owned();
        // Find min and max along each dimension using ndarray operations
        let scaled_min: Array1<f64> = scaled_training_points.map_axis(Axis(0), |col| {
            col.iter().copied().fold(f64::INFINITY, f64::min)
        });
        let scaled_max: Array1<f64> = scaled_training_points.map_axis(Axis(0), |col| {
            col.iter().copied().fold(f64::NEG_INFINITY, f64::max)
        });
        // Estimate sparsity
        let sparsity: f64 = (&scaled_max - &scaled_min).mean().unwrap();
        // Create hash for training data
        let points_for_hash: Vec<Vec<f64>> = scaled_training_points
            .axis_iter(Axis(0))
            .map(|row| row.to_vec())
            .collect();
        let training_hash = SpatialHash::build(&points_for_hash, 1.0);
        // Train kernel
        let training_kernel = match sparsity {
            sparsity if sparsity > 2. => {
                kernel.hashed_kernel_construction(
                    scaled_training_points.view(),
                    training_error.view(),
                    &training_hash,
                )
            }
            _ => {
                kernel.naive_kernel_construction(
                    scaled_training_points.view(),
                    training_error.view(),
                )
            }
        };
        // Convert to csr
        let training_kernel = training_kernel.to_csr();
        // Construct the LDL decomposition
        let numeric = LdlNumeric::new(training_kernel.view())
            .expect("Numeric LDL analysis failed");
        // Evaluate the alphas
        let alpha_vec = numeric.solve(training_values.as_slice().unwrap());
        let alpha = Array1::from_vec(alpha_vec);

        let mut gp = GP {
            ndim,
            kernel,
            training_points: scaled_training_points,
            training_values,
            training_error,
            training_kernel,
            training_hash,
            alpha,
            lnL: 0.,
            ldl_numeric: Some(numeric),
        };
        // Evaluate lnL
        gp.log_marginal_likelihood();
        // Return GP
        gp
    }

    /// Provide a way to recompose the LDL matrix after loading
    pub fn generate_ldl(&mut self) {
        // Construct the LDL decomposition
        self.ldl_numeric = Some(LdlNumeric::new(self.training_kernel.view())
            .expect("Numeric LDL analysis failed"));
    }

    pub fn retrain(
        &mut self,
        scale: ArrayView1<f64>,
        whitenoise: f64,
        msq_order: i32,
    ) {
        /*!
        Retrain the Gaussian Process struct with different hyperparameters
        */
        // Initialize kernel 
        let kernel = Kernel::new(self.ndim, msq_order, scale, whitenoise);
        // Rescale training_points
        let training_points = &self.training_points * (&self.kernel.scale() / &scale);
        // Find min and max along each dimension using ndarray operations
        let scaled_min: Array1<f64> = training_points.map_axis(Axis(0), |col| {
            col.iter().copied().fold(f64::INFINITY, f64::min)
        });
        let scaled_max: Array1<f64> = training_points.map_axis(Axis(0), |col| {
            col.iter().copied().fold(f64::NEG_INFINITY, f64::max)
        });
        // Estimate sparsity
        let sparsity: f64 = (&scaled_max - &scaled_min).mean().unwrap();
        // Create hash for training data
        let points_for_hash: Vec<Vec<f64>> = training_points
            .axis_iter(Axis(0))
            .map(|row| row.to_vec())
            .collect();
        let training_hash = SpatialHash::build(&points_for_hash, 1.0);
        // Train kernel
        let training_kernel = match sparsity {
            sparsity if sparsity > 2. => {
                kernel.hashed_kernel_construction(
                    training_points.view(),
                    self.training_error.view(),
                    &training_hash,
                )
            }
            _ => {
                kernel.naive_kernel_construction(
                    training_points.view(),
                    self.training_error.view(),
                )
            }
        };
        // Convert to csr
        let training_kernel = training_kernel.to_csr();
        self.training_points = training_points;
        self.training_hash = training_hash;
        self.kernel = kernel;
        self.training_kernel = training_kernel;
        // ---- numeric analysis ----
        let numeric = LdlNumeric::new(self.training_kernel.view())
            .expect("Numeric LDL analysis failed");
        let alpha_vec = numeric.solve(self.training_values.as_slice().unwrap());
        let alpha = Array1::from_vec(alpha_vec);

        self.ldl_numeric = Some(numeric);
        self.alpha = alpha;
        self.log_marginal_likelihood();
    }

    /// Loop through dense kernel entries and apply a specified function
    pub fn for_each_kernel_entry<F: FnMut(usize, f64)>(
        &self,
        eval_point: ArrayView1<f64>,
        mut f: F,
    ) {
        debug_assert_eq!(eval_point.len(), self.ndim);
        // Scale eval point
        let eval_point_scaled = self.kernel.scale_sgl(eval_point);
        // visit only nearby training points
        self.training_hash.for_each_neighbor(
            eval_point_scaled.as_slice().unwrap(), |i| {
                // Identify training point
                let xi = self.training_points.row(i);
                // Calculate distance using iterator
                let r2: f64 = eval_point_scaled.iter()
                    .zip(xi.iter())
                    .map(|(ei, xi)| (ei-xi).powi(2))
                    .sum();

                // Check if point is blank
                if r2 >= 1.0 { return; }
                // Get r
                let r = r2.sqrt();
                // Calculate kernel
                let kij = self.kernel.eval_sgl(r, r2);
                f(i,kij);
            }
        );
    }

    fn fill_kernel_vec(&self, eval_point: ArrayView1<f64>, out: &mut [f64]) {
        /*! Create a vector k(x, x') for an individual point x'
        */
        debug_assert_eq!(out.len(), self.training_points.nrows());
        debug_assert_eq!(eval_point.len(), self.ndim);
        out.fill(0.);
        self.for_each_kernel_entry(eval_point,
            |i, kij| {out[i] = kij;});
    }

    pub fn predict_mean_sgl(&self, eval_point: ArrayView1<f64>) -> f64 {
        /*!
        Predict the mean value for the realizations of the Gaussian Process
        at a single point
        */
        // Scale evaluation samples
        let eval_point_scaled = self.kernel.scale_sgl(eval_point);
        // initialize mean
        let mut mean: f64 = 0.;
        // Only evaluate close points
        self.training_hash.for_each_neighbor(
            eval_point_scaled.as_slice().unwrap(), |i| {
                // Identify the point
                let xi = self.training_points.row(i);
                // Calculate r2 using iterator
                let r2: f64 = eval_point_scaled.iter()
                    .zip(xi.iter())
                    .map(|(ei, xi)| (ei-xi).powi(2))
                    .sum();
                // Check for zeros
                // iff sqrt(r2) > 1 then r2 > 1
                if r2 >= 1.0 { return; }
                // take the square root
                let r = r2.sqrt();
                // Identify the kernel value
                let val = self.kernel.eval_sgl(r, r2);
                mean += val * self.alpha[i];
            }
        );
        mean
    }
    pub fn predict_mean(&self, eval_points: ArrayView2<f64>) -> Array1<f64> {
        /*!
        Predict the mean value for the realizations of the Gaussian Process
        at a vector of points.
        */
        // Scale all points at once
        let scaled_points = self.kernel.scale_arr(eval_points);

        // Allocate output
        let neval = eval_points.nrows();
        let mut means = vec![0.; neval];
        
        // Process each point
        for (idx, scaled_row) in scaled_points.outer_iter().enumerate() {
            let mut mean = 0.;
            // Get slice once
            let scaled_slice = scaled_row.as_slice().unwrap();
            // Loop spatial hash
            self.training_hash.for_each_neighbor(scaled_slice, |i| {
                // Identify training point
                let xi = self.training_points.row(i);
                // Calculate r2 using iterator
                let mut r2 = 0.;
                for (ei, xi) in scaled_slice.iter()
                    .zip(xi.iter()) {
                    let diff = ei - xi;
                    r2 += diff * diff;
                }
                // Check for zeros
                if r2 < 1.0 {
                    let r = r2.sqrt();
                    mean += self.kernel.eval_sgl(r, r2) * self.alpha[i];
                }
            });
            means[idx] = mean;
        }
        Array1::from_vec(means)
    }

    pub fn predict_var(&self, eval_points: ArrayView2<f64>) -> Array1<f64> {
        // Obtain reference to numeric
        let numeric = self.ldl_numeric
            .as_ref()
            .expect("GP not trained; call train method first!");
        // Identify number of evaluation points
        let ntrain = self.training_points.nrows();
        let mut k_eval_pt = vec![0.;ntrain];
        // Loop evaluation points
        eval_points.axis_iter(Axis(0)).map(|x| {
            // Build k(x,x')
            self.fill_kernel_vec(x, &mut k_eval_pt);
            // Solve for q
            let q = numeric.solve(&k_eval_pt);
            // k(x,x')^T * q
            let quad: f64 = k_eval_pt.iter()
                .zip(q.iter())
                .map(|(k, q)| k * q)
                .sum();
            // k(x, x) = 1 for Wendland
            let variance = 1.0 - quad;
            // numerical safety
            variance.max(0.)
        })
        .collect()
    }

    /// Predict the covariance between two points
    pub fn predict_cov_sgl(
        &self,
        x1: ArrayView1<f64>,
        x2: ArrayView1<f64>,
    ) -> f64 {
        // Need trained GP
        let numeric = self.ldl_numeric.as_ref()
            .expect("GP not trained; call generate_ldl() first!");
        // Identify number of training points
        let ntrain = self.training_points.nrows();
        // Build k(x1, X) and k(x2, X)
        let mut k1 = vec![0.; ntrain];
        let mut k2 = vec![0.; ntrain];
        // Fill the kernel vectors
        self.fill_kernel_vec(x1, &mut k1);
        self.fill_kernel_vec(x2, &mut k2);
        // Solve K q = k2
        let q = numeric.solve(&k2);
        // Compute k1^T q
        let quad: f64 = k1.iter()
            .zip(q.iter())
            .map(|(k,q)| k * q)
            .sum();
        // Direct kernel k(x1, x2)
        // Scale inputs
        let x1s = self.kernel.scale_sgl(x1);
        let x2s = self.kernel.scale_sgl(x2);
        // Calculate distance
        let r2: f64 = x1s.iter()
            .zip(x2s.iter())
            .map(|(k1, k2)| (k1 - k2).powi(2))
            .sum();

        // Evaluate kernel
        let k12 = if r2 > 1.0 { 0. } else {
            let r = r2.sqrt();
            self.kernel.eval_sgl(r, r2)
        };
        // Evaluate covariance
        let cov = k12 - quad;
        // Numerical safety
        if cov.abs() < 1e-14 { 0. } else { cov }
    }
    /// Predict the variance at a single point
    pub fn predict_var_sgl(&self, eval_point: ArrayView1<f64>) -> f64 {
        let eval_points = eval_point.insert_axis(Axis(0));
        self.predict_var(eval_points)[0]
    }

    /// Save the Gaussian Process to a specified file and group
    pub fn save_to_hdf5(&self, filename: &str, group_name: &str) -> H5Result<()> {
        // Open or create the file
        let file = if std::path::Path::new(filename).exists() {
            File::open_rw(filename)? // open for read/write (existing)
        } else {
            File::create(filename)?
        };
        // Define group to point at intended path
        let group = if file.link_exists(group_name) {
            //file.group(group_name)?
            panic!("Cannot save new GP! {group_name} already taken!");
        } else {
            file.create_group(group_name)?
        };
        // Save basic scalars as attributes
        group.new_attr::<f64>().create("whitenoise")?
            .write_scalar(&self.kernel.whitenoise())?;
        group.new_attr::<usize>().create("ndim")?
            .write_scalar(&self.ndim)?;
        group.new_attr::<i32>().create("msq_order")?
            .write_scalar(&self.kernel.msq_order())?;
        // Version
        group.new_attr::<u32>().create("format_version")?
            .write_scalar(&1)?;
        // lnL
        group.new_attr::<f64>().create("lnL")?
            .write_scalar(&self.lnL)?;

        // Datasets
        // scale is a 1D vector, so that's easy
        group.new_dataset_builder()
            .with_data(self.kernel.scale().as_slice().unwrap())
            .create("scale")?;
        // training_points is 2D
        // TODO I know there has got to be a better way to do this
        // but I don't know the better way to do this
        let pts_flat: Vec<f64> = self.training_points
            .iter().copied().collect();
        // Save training values
        group.new_dataset_builder()
            .with_data(&pts_flat)
            .create("training_points")?;

        // Training values is also 1-D, so easy
        group.new_dataset_builder()
            .with_data(self.training_values.as_slice().unwrap())
            .create("training_values")?;

        // Training error is also 1-D, so easy
        group.new_dataset_builder()
            .with_data(self.training_error.as_slice().unwrap())
            .create("training_error")?;

        // alpha is also a 1-D array, so easy
        group.new_dataset_builder()
            .with_data(self.alpha.as_slice().unwrap())
            .create("alpha")?;
        // finally, extract csr data
        let csr = &self.training_kernel;
        group.new_dataset_builder()
            .with_data(csr.indices())
            .create("kernel_indices")?;
        group.new_dataset_builder()
            .with_data(csr.data())
            .create("kernel_data")?;
        group.new_dataset_builder()
            .with_data(csr.indptr().raw_storage())
            .create("kernel_indptr")?;
        Ok(())
    }
    /// Load a Gaussian Process previously saved in HDF5
    /// at the given file and group
    pub fn load_from_hdf5(filename: &str, group_name: &str) -> H5Result<GP> {
        // Try to read the file
        let file = File::open(filename)?;
        let group = file.group(group_name)?;

        // Read scalars
        let ndim:           usize = group.attr("ndim")?.read_scalar()?;
        let msq_order:      i32 = group.attr("msq_order")?.read_scalar()?;
        let whitenoise:     f64 = group.attr("whitenoise")?.read_scalar()?;
        #[allow(non_snake_case)]
        let lnL:            f64 = group.attr("lnL")?.read_scalar()?;
        // ---- read datasets ----
        let scale_vec: Vec<f64> = group.dataset("scale")?
            .read_raw()?
            .to_vec();
        let scale = Array1::from_vec(scale_vec);
        // Construct Kernel parameters
        let kernel = Kernel::new(ndim, msq_order, scale.view(), whitenoise);

        // Un-flatten training points
        let pts_flat: Vec<f64> = group.dataset("training_points")?
            .read_raw()?
            .to_vec();
        let ntrain = pts_flat.len() / ndim;
        let training_points = Array2::from_shape_vec((ntrain, ndim),pts_flat)
            .expect("Failed to reshape training points");

        // Carry on with other 1D arrays
        let training_values_vec: Vec<f64> =
            group.dataset("training_values")?
            .read_raw()?
            .to_vec();
        let training_values = Array1::from_vec(training_values_vec);
        let training_error_vec: Vec<f64> =
            group.dataset("training_error")?
            .read_raw()?
            .to_vec();
        let training_error = Array1::from_vec(training_error_vec);
        // alpha
        let alpha_vec: Vec<f64> =
            group.dataset("alpha")?
            .read_raw()?
            .to_vec();
        let alpha = Array1::from_vec(alpha_vec);
        // Reconstruct CSR
        let indptr: Vec<usize> = group.dataset("kernel_indptr")?
            .read_raw()?
            .to_vec();
        let indices: Vec<usize> = group.dataset("kernel_indices")?
            .read_raw()?
            .to_vec();
        let data: Vec<f64> = group.dataset("kernel_data")?
            .read_raw()?
            .to_vec();
        // Instantiate kernel
        let training_kernel = CsMat::new((ntrain,ntrain), indptr, indices, data);
        // Train hash
        let points_for_hash: Vec<Vec<f64>> = training_points
            .axis_iter(Axis(0))
            .map(|row| row.to_vec())
            .collect();
        let training_hash = SpatialHash::build(&points_for_hash, 1.0);
        Ok(GP {
            ndim,
            kernel,
            training_points,
            training_values,
            training_error,
            training_kernel,
            training_hash,
            alpha,
            lnL,
            ldl_numeric: None,
        })
    }
    /// Evaluate the log marginal likelihood of 
    /// the Gaussian process hyperparameters
    pub fn log_marginal_likelihood(&mut self) -> f64 {
        // Hopefully the compiler optimizes this
        let ln2pi: f64 = (2. * std::f64::consts::PI).ln();
        // get references to numeric and alpha
        let numeric = self.ldl_numeric.as_ref().unwrap();
        // Sum in quadrature
        let quad: f64 = self.training_values.iter()
            .zip(self.alpha.iter())
            .map(|(y, a)| y * a)
            .sum();
        // Get diagonal
        let d = numeric.d();
        // Initialize the log of the determinant
        let logdet: f64 = d.iter().map(|&di| di.ln()).sum();
        // Return the log likelihood
        self.lnL = (-0.5 * quad) - (0.5 * logdet) - 
            (0.5 * (self.alpha.len() as f64) * ln2pi);
        // Quantify it
        self.lnL
    }
    /// Draw a realization of the Gaussian Process
    pub fn draw_realization<'a>(
        &'a self,
        rng: &mut PCG64Stream,
    ) -> GPRealization<'a> {
        // Get reference to numeric
        let numeric = self.ldl_numeric
            .as_ref()
            .expect("GP not trained");
        // Get number of training points
        let ntrain = self.training_points.nrows();

        // 1. Get a standard normal
        let mut z = vec![0.;ntrain];
        rng.fill_standard_normal(&mut z);
        let z0 = z.clone();
        
        // 2. v = sqrt(D) * z
        let d = numeric.d();
        let v: Vec<f64> = z.iter()
            .zip(d.iter())
            .map(|(&zi, &di)| di.sqrt() * zi)
            .collect();

        // 3. f0 = L * v
        let l = numeric.l();
        let mut f0 = vec![0.;ntrain];
        // because L has implicit unit diagonal:
        //
        // f0[i] = v[i] + sum{j<i} L_ij v[j]
        //
        if l.is_csr() {
            // rows
            for (i, row) in l.outer_iterator().enumerate() {
                let acc = v[i] + row.iter()
                    .filter_map(|(j,lij)| (j < i).then(|| lij * v[j]))
                    .sum::<f64>();
                f0[i] = acc;
            }
        } else if l.is_csc() {
            // columns: accumulate forward
            for (j, col) in l.outer_iterator().enumerate() {
                let vj = v[j];
                for (i, lij) in col.iter().filter(|&(i, _)| i > j) {
                    f0[i] += lij * vj;
                }
            }
        }
        // 4. Solve for u
        let u = numeric.solve(&f0);
        // Identify f_train
        // 5. Compute K * u
        // ftrain = y + f0 - K u
        let k = &self.training_kernel;
        let mut ku = vec![0.;ntrain];
        // More matrix math
        if k.is_csr() {
            // ku[i] = sum_j K_ij * u[j]
            for (i, row) in k.outer_iterator().enumerate() {
                ku[i] = row.iter()
                    .map(|(j,kij)| kij * u[j])
                    .sum();
            }
        } else if k.is_csc() {
            // ku[i] += K_ij * u[j]
            for (j, col) in k.outer_iterator().enumerate() {
                let uj = u[j];
                for (i, kij) in col.iter() {
                    ku[i] += kij * uj;    
                }
            }
        }
        // 6. f_train = y + f0 - K*u
        let f_train: Vec<f64> = self.training_values.iter()
            .zip(f0.iter().zip(ku.iter()))
            .map(|(y, (f, k))| y + f - k)
            .collect();
        // Precompute beta
        let beta = numeric.solve(&f_train);
        GPRealization::new(self, f_train, beta, z0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::arr1;

    fn vec2d_to_array2(v: &[Vec<f64>]) -> Array2<f64> {
        let nrows = v.len();
        let ncols = v[0].len();
        let flat: Vec<f64> = v.iter().flat_map(
            |row| row.iter().copied()).collect();
        Array2::from_shape_vec((nrows,ncols), flat).unwrap()
    }
    
    #[test]
    fn gp_interpolates_training_points() {
        // define training points
        let x_train = vec2d_to_array2(&[vec![0.], vec![0.5], vec![1.0]]);
        // define training values
        let y_train = arr1(&[0., 1., 0.]);
        // Define training error
        let y_train_err = arr1(&[0., 0., 0.]);
        let kernel_noise = 0.;
        // Define scale
        let scale = arr1(&[1.0]);
        // Loop order
        for order in 0..4 {

            // Initialize gp
            let gp = GP::train(
                x_train.view(),
                y_train.view(),
                y_train_err.view(),
                scale.view(),
                kernel_noise,
                order,
            );

            // Evaluate prediction
            let prediction = gp.predict_mean(x_train.view());

            // Check values
            for i in 0..x_train.nrows() {
                assert!(
                    (prediction[i] - y_train[i]).abs() < 1e-6,
                    "prediction[i] = {}, y_train = {}", prediction[i], y_train[i],
                );
            }
        }
    }

    #[test]
    fn compact_support_gives_zero_mean_far_away() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5]]);
        // Define training values
        let y_train = arr1(&[8., -6.]);
        let y_train_err = arr1(&[0.04,0.08]);
        let whitenoise = 1e-8;
        let scale = arr1(&[1.]);
        //loop order
        for order in 0..4 {
            // Initialize gp
            let gp = GP::train(
                x_train.view(),
                y_train.view(),
                y_train_err.view(),
                scale.view(),
                whitenoise,
                order,
            );
            // Test far away
            let far = arr1(&[10.0]);
            let prediction = gp.predict_mean_sgl(far.view());
            assert!(prediction.abs() < 1e-10);
        }
    }

    #[test]
    fn test_interpolation() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![1.]]);
        // Define training values
        let y_train = arr1(&[0., 1.]);
        let y_train_err = arr1(&[0.04,0.08]);
        let whitenoise = 1e-8;
        let scale = arr1(&[1.]);
        //loop order
        for order in 0..4 {
            // Iniitialize gp
            let gp = GP::train(
                x_train.view(),
                y_train.view(),
                y_train_err.view(),
                scale.view(),
                whitenoise, 
                order, 
            );
            // Test between
            let between = arr1(&[0.5]);
            let prediction = gp.predict_mean_sgl(between.view());
            assert!(prediction.abs() > 0.);
            assert!(prediction.abs() < 1.);
        }
    }
    #[test]
    fn disconnected_components_do_not_interact() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.2],vec![5.0],vec![5.2]]);
        // Define training values
        let y_train = arr1(&[1.,1.,-1.,-1.]);
        let err = arr1(&[0.,0.,0.,0.,]);
        let whitenoise = 1e-8;
        let scale = arr1(&[1.0]);

        for order in 0..4 {
            // Initialize GP
            let gp = GP::train(
                x_train.view(),
                y_train.view(),
                err.view(),
                scale.view(),
                whitenoise,
                order,
            );
            // Evaluate near each block
            let near_left = arr1(&[0.1]);
            let near_right = arr1(&[5.1]);
            let prediction_left = gp.predict_mean_sgl(near_left.view());
            let prediction_right = gp.predict_mean_sgl(near_right.view());
            // Check that values make some kind of sense
            assert!(prediction_left > 0.5);
            assert!(prediction_right < -0.5);
        }
    }
    #[test]
    fn variance_zero_at_training_points() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.],]);
        let y_train = arr1(&[0.,1.,0.]);
        let y_err = arr1(&[0.;3]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.]);
        // Loop order
        for order in 0..4 {
            // train the gp
            let gp = GP::train(
                x_train.view(), y_train.view(), y_err.view(),
                scale.view(), whitenoise, order
            );
            // Find variance
            let vars = gp.predict_var(x_train.view());
            // Loop points
            for (i, var) in vars.iter().enumerate() {
                assert!(
                    *var < 1e-6,
                    "variance at training point {i} too large: {var}"
                );
            }
        }
    }
    #[test]
    fn variance_reverts_to_prior_far_away() {
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.3]]);
        let y_train = arr1(&[1.,-1.]);
        let y_err = arr1(&[0.,0.]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.]);
        // Loop order
        for order in 0..4 {
            // train the gp
            let gp = GP::train(
                x_train.view(), y_train.view(), y_err.view(),
                scale.view(), whitenoise, order
            );
            // Find a far away point
            let far = arr1(&[10.]);
            let var = gp.predict_var_sgl(far.view());
            // Check that variance is close to 1.
            assert!((var - 1.).abs() < 1e-6, "far variance wrong: {var}");
        }
    }
    #[test]
    fn variance_reduced_inside_support() {
        let x_train = vec2d_to_array2(&[vec![0.],vec![1.]]);
        let y_train = arr1(&[0.,1.]);
        let y_err = arr1(&[0.;2]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.]);
        // Loop order
        for order in 0..4 {
            // train the gp
            let gp = GP::train(
                x_train.view(), y_train.view(), y_err.view(),
                scale.view(), whitenoise, order
            );
            // Find the midpoint
            let mid = arr1(&[0.5]);
            let var = gp.predict_var_sgl(mid.view());
            assert!(var > 0., "variance should be positive");
            assert!(var < 1., "variance should be reduced below prior");
        }
    }
    #[test]
    fn predictive_covariance_is_symmetric() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let y_err = arr1(&[0.; 3]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), y_err.view(),
            scale.view(), whitenoise, 2
        );
        // Define some sample locations
        let x1 = arr1(&[0.2]);
        let x2 = arr1(&[0.7]);
        // Get covariance
        let c12 = gp.predict_cov_sgl(x1.view(), x2.view());
        let c21 = gp.predict_cov_sgl(x2.view(), x1.view());
        assert!((c12-c21).abs() < 1e-10, "cov not symmetric: {c12} vs {c21}");

    }
    #[test]
    fn predictive_covariance_reduces_to_variance() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let y_err = arr1(&[0.,0.,0.]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), y_err.view(),
            scale.view(), whitenoise, 2
        );
        // Define some sample locations
        let x = arr1(&[0.3]);
        let var = gp.predict_var_sgl(x.view());
        let cov = gp.predict_cov_sgl(x.view(), x.view());
        assert!((var - cov).abs() <1e-8, "var {var} vs cov {cov}");
    }
    #[test]
    fn predictive_covariance_zero_far_apart() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.3]]);
        let y_train = arr1(&[-1.0, 1.0]);
        let y_err = arr1(&[0.,0.]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), y_err.view(),
            scale.view(), whitenoise, 2);
        // Define some sample locations
        let x1 = arr1(&[0.1]);
        let x2 = arr1(&[10.0]);
        // Get covariance
        let c12 = gp.predict_cov_sgl(x1.view(), x2.view());
        let c21 = gp.predict_cov_sgl(x2.view(), x1.view());
        assert!(c12.abs() < 1e-10, "far cov (c12) not zero: {c12}");
        assert!(c21.abs() < 1e-10, "far cov (c21) not zero: {c21}");
    }
    #[test]
    fn log_marginal_likelihood_is_finite() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let y_err = arr1(&[0.; 3]);
        let whitenoise = 1e-6;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), y_err.view(), 
            scale.view(), whitenoise, 2);
        // Define some sample locations
        let ll = gp.lnL;
        assert!(ll.is_finite(), "log likelihood not finite");
        assert!(ll < 0., "log likelihood should not be negative, got {ll}")
    }
    #[test]
    fn log_marginal_prefers_smooth_data() {
        // Define training data
        let x_train = vec2d_to_array2(
            &[vec![0.],vec![0.25],vec![0.5],vec![0.75],vec![1.0]]
        );
        let y_train_smooth = arr1(&[0., 0.7, 1., 0.7, 0.]);
        let y_train_noisy = arr1(&[0.3, -1.2, 0.4, 2.0, -0.7]);
        let err = arr1(&[0.; 5]);
        let whitenoise = 1e-6;
        let scale = arr1(&[1.0]);
        // Build GP smooth
        let gp_smooth = GP::train(
            x_train.view(),
            y_train_smooth.view(),
            err.view(),
            scale.view(),
            whitenoise,
            2,
        );
        // Build GP noisy
        let gp_noisy = GP::train(
            x_train.view(),
            y_train_noisy.view(),
            err.view(),
            scale.view(),
            whitenoise,
            2,
        );
        // Get log likelihoods
        let ll_smooth = gp_smooth.lnL;
        let ll_noisy = gp_noisy.lnL;
        // check that the smooth data is better
        assert!(ll_smooth > ll_noisy,
            "smooth ll {ll_smooth} not > noisy ll {ll_noisy}");
    }
    #[test]
    fn predictive_covariance_psd_2x2() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let err = arr1(&[0.; 3]);
        let whitenoise = 1e-6;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), err.view(), 
            scale.view(), whitenoise, 2);
        // Define some sample locations
        let x1 = arr1(&[0.25]);
        let x2 = arr1(&[0.75]);
        // Identify vars
        let v1 = gp.predict_var_sgl(x1.view());
        let v2 = gp.predict_var_sgl(x2.view());
        let c12 = gp.predict_cov_sgl(x1.view(), x2.view());
        // Get determinant
        let det = v1 * v2 - c12 * c12;
        assert!(det >= -1e-10, "covariance not PSD, det={det}");
    }
    #[test]
    fn posterior_variance_bounded_by_prior() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.4],vec![0.8]]);
        let y_train = arr1(&[0.2,-0.5,0.1]);
        let err = arr1(&[0.;3]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), err.view(), 
            scale.view(), whitenoise, 2);
        for i in 0..20{
            let x = arr1(&[i as f64 / 20.]);
            let var = gp.predict_var_sgl(x.view());
            assert!(var >= -1e-12, "negative variance {var}");
            assert!(var <= 1. + 1e-12, "variance exceeds prior {var}");
        }
    }
    #[test]
    fn log_marginal_invariant_under_permutation() {
        // Get two sets of training data
        let x1 = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y1 = arr1(&[0.,1.,0.]);
        let x2 = vec2d_to_array2(&[vec![1.],vec![0.],vec![0.5]]);
        let y2 = arr1(&[0.,0.,1.]);
        let err = arr1(&[0.;3]);
        let whitenoise = 1e-6;
        let scale = arr1(&[1.0]);
        // train gps
        let gp1 = GP::train(
            x1.view(), y1.view(), err.view(), 
            scale.view(), whitenoise, 2);
        let gp2 = GP::train(
            x2.view(), y2.view(), err.view(), 
            scale.view(), whitenoise, 2);
        // Get the log likelihood
        let lnl_1 = gp1.lnL;
        let lnl_2 = gp2.lnL;
        assert!((lnl_1 - lnl_2).abs() < 1e-8,
            "log likelihood not permutation invariant");
    }
    #[test]
    fn more_noise_means_worse_likelihood() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let err = arr1(&[0.; 3]);
        let scale = arr1(&[1.0]);
        let whitenoise_lo = 1e-8;
        let whitenoise_hi = 1e-1;
        // Build GP_lo
        let mut gp = GP::train(
            x_train.view(), y_train.view(), err.view(), 
            scale.view(), whitenoise_lo, 2);
        // Get the log likelihoods
        let lnl_lo = gp.lnL;
        // Build GP_hi
        gp.retrain(scale.view(),whitenoise_hi,2);
        // Get the log likelihoods
        let lnl_hi = gp.lnL;
        // More noise = worse likelihood
        assert!(lnl_lo > lnl_hi, "More noise, but better fit?");
    }
    #[test]
    fn retrained_gp_has_correct_mean() {
        // Define training data
        let x_train = vec2d_to_array2(&[
            vec![0.],vec![0.25],vec![0.5],vec![0.75],vec![1.0]
        ]);
        let y_train = arr1(&[0., 1., 2., 3., 4.]);
        let y_err = arr1(&[0.;5]);
        let whitenoise = 1e-10;
        let msq_order: i32 = 1;
        let scale_lo = arr1(&[1.]);
        let scale_hi = arr1(&[2.]);
        let mut gp = GP::train(
            x_train.view(),
            y_train.view(),
            y_err.view(),
            scale_lo.view(),
            whitenoise,
            msq_order,
        );
        // Get sample locations
        let x_sample = vec2d_to_array2(&[
            vec![0.0],vec![0.1],vec![0.2],vec![0.3],vec![0.4],
            vec![0.5],vec![0.6],vec![0.7],vec![0.8],vec![0.9],
        ]);
        // Evaluate mean
        let y_lo = gp.predict_mean(x_sample.view());
        // Retrain the model
        gp.retrain(scale_hi.view(), whitenoise, msq_order);
        // Evaluate mean
        let y_hi = gp.predict_mean(x_sample.view());
        for (x_lo, x_hi) in y_lo.iter().zip(y_hi.iter()) {
            assert![(x_hi - x_lo).abs() < 0.3];
        }
    }
    #[test]
    fn realizations_have_correct_mean() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let err = arr1(&[0.; 3]);
        let whitenoise = 1e-6;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), err.view(), 
            scale.view(), whitenoise, 2);
        // Define some sample locations
        let x = vec2d_to_array2(&[vec![0.1], vec![0.2], vec![0.3]]);
        let mean = gp.predict_mean(x.view());
        // Define a random number generator
        let mut rng = PCG64Stream::new(1234, 5678);
        // do some tests with it
        let mut acc = arr1(&vec![0.;x.nrows()]);
        let n = 2000;
        for _ in 0..n {
            // Get a realization
            let r = gp.draw_realization(&mut rng);
            for i in 0..x.nrows() {
                acc[i] += r.value(x.row(i));
            }
        }
        let est: Vec<f64> = acc.iter().map(
            |a| a / n as f64
        ).collect();
        for i in 0..acc.len() {
            assert!((est[i] - mean[i]).abs() < 5e-2,
                "realization mean mismatch: {} vs {}",
                est[i], mean[i]);
        }
    }
    #[test]
    fn realization_deterministic_at_training_points() {
        // Define training data
        let x_train = vec2d_to_array2(&[vec![0.],vec![0.5],vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let err = arr1(&[0.; 3]);
        let whitenoise = 1e-12;
        let scale = arr1(&[1.0]);
        // Build GP
        let gp = GP::train(
            x_train.view(), y_train.view(), err.view(), 
            scale.view(), whitenoise, 2);
        // Get a random number generator
        let mut rng = PCG64Stream::new(42,99);
        // keep track of accuracy
        let mut acc = arr1(&vec![0.; x_train.nrows()]);
        let n = 5000;
        // Evaluate 5000 realziations of each training point
        for _ in 0..n {
            let real = gp.draw_realization(&mut rng);
                for i in 0..x_train.nrows() {
                    acc[i] += real.value(x_train.row(i));
                }
        }
        // Divide by number of training points
        for i in 0..x_train.nrows() {
            let est: f64 = acc[i] / n as f64;
            assert!((est - y_train[i]).abs() < 5e-2);
        }
    }
    #[test]
    fn realization_training_point_variance_collapses() {
        let x_train = vec2d_to_array2(&[vec![0.], vec![0.5], vec![1.0]]);
        let y_train = arr1(&[0., 1., 0.]);
        let err = arr1(&[0.; 3]);
        let whitenoise = 1e-10;
        let scale = arr1(&[1.0]);
        let gp = GP::train(
            x_train.view(), y_train.view(), err.view(), 
            scale.view(), whitenoise, 2);

        let mut rng = PCG64Stream::new(7, 11);

        let n = 4000;
        let m = x_train.len();

        let mut sum = arr1(&vec![0.; m]);
        let mut sum2 = arr1(&vec![0.; m]);

        for _ in 0..n {
            let r = gp.draw_realization(&mut rng);
            for i in 0..m {
                let v = r.value(x_train.row(i));
                sum[i] += v;
                sum2[i] += v * v;
            }
        }

        for i in 0..m {
            let mean : f64 = sum[i] / n as f64;
            let var : f64 = sum2[i] / n as f64 - mean * mean;

            assert!(
                (mean - y_train[i]).abs() < 5e-2,
                "mean mismatch at training point {}: {} vs {}",
                i, mean, y_train[i]
            );
            assert!(
                var < 5e-3,
                "variance did not collapse at training point {}: {}",
                i, var
            );
        }
    }
}