sklears_kernel_approximation/sparse_gp/

inference.rs

//! Scalable inference methods for sparse Gaussian Processes
//!
//! This module implements various scalable inference algorithms including
//! direct matrix inversion, Preconditioned Conjugate Gradient (PCG),
//! and Lanczos eigendecomposition methods.

use crate::sparse_gp::core::*;
use crate::sparse_gp::kernels::{KernelOps, SparseKernel};

use scirs2_core::ndarray::s;
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::essentials::Uniform as RandUniform;
use scirs2_core::random::thread_rng;
use scirs2_linalg::compat::ArrayLinalgExt;
use sklears_core::error::{Result, SklearsError};

/// Scalable inference method implementations
pub struct ScalableInference;

impl ScalableInference {
    /// Perform scalable prediction using the specified method
    pub fn predict<K: SparseKernel>(
        method: &ScalableInferenceMethod,
        k_star_m: &Array2<f64>,
        inducing_points: &Array2<f64>,
        alpha: &Array1<f64>,
        kernel: &K,
        noise_variance: f64,
    ) -> Result<Array1<f64>> {
        match method {
            ScalableInferenceMethod::Direct => Self::predict_direct(k_star_m, alpha),
            ScalableInferenceMethod::PreconditionedCG {
                max_iter,
                tol,
                preconditioner,
            } => Self::predict_with_pcg(
                k_star_m,
                inducing_points,
                kernel,
                noise_variance,
                *max_iter,
                *tol,
                preconditioner,
            ),
            ScalableInferenceMethod::Lanczos { num_vectors, tol } => Self::predict_with_lanczos(
                k_star_m,
                inducing_points,
                kernel,
                noise_variance,
                *num_vectors,
                *tol,
            ),
        }
    }

    /// Direct prediction using precomputed alpha
    fn predict_direct(k_star_m: &Array2<f64>, alpha: &Array1<f64>) -> Result<Array1<f64>> {
        Ok(k_star_m.dot(alpha))
    }

    /// Prediction using Preconditioned Conjugate Gradient
    fn predict_with_pcg<K: SparseKernel>(
        k_star_m: &Array2<f64>,
        inducing_points: &Array2<f64>,
        kernel: &K,
        noise_variance: f64,
        max_iter: usize,
        tol: f64,
        preconditioner: &PreconditionerType,
    ) -> Result<Array1<f64>> {
        let m = inducing_points.nrows();

        // Reconstruct the system matrix A = K_mm + noise
        let mut a_matrix = kernel.kernel_matrix(inducing_points, inducing_points);
        for i in 0..m {
            a_matrix[(i, i)] += noise_variance;
        }

        // Right-hand side for prediction
        let rhs = k_star_m.t().dot(&Array1::ones(k_star_m.nrows()));

        // Solve A * x = rhs using PCG
        let solution = PreconditionedCG::solve(&a_matrix, &rhs, max_iter, tol, preconditioner)?;

        Ok(k_star_m.dot(&solution))
    }

    /// Prediction using Lanczos method
    fn predict_with_lanczos<K: SparseKernel>(
        k_star_m: &Array2<f64>,
        inducing_points: &Array2<f64>,
        kernel: &K,
        noise_variance: f64,
        num_vectors: usize,
        tol: f64,
    ) -> Result<Array1<f64>> {
        // Reconstruct kernel matrix
        let k_mm = kernel.kernel_matrix(inducing_points, inducing_points);

        // Apply Lanczos algorithm for eigendecomposition
        let (eigenvals, eigenvecs) = LanczosMethod::eigendecomposition(&k_mm, num_vectors, tol)?;

        // Use eigendecomposition for prediction
        // This is a simplified version - full implementation would use proper alpha reconstruction
        let k_star_transformed = k_star_m.dot(&eigenvecs);

        // Apply eigenvalue scaling (simplified)
        let scaled_eigenvals = eigenvals.mapv(|x| {
            if x > 1e-10 {
                1.0 / (x + noise_variance)
            } else {
                0.0
            }
        });

        // Compute prediction (simplified)
        let prediction = k_star_transformed.dot(&scaled_eigenvals);
        Ok(prediction)
    }
}

/// Preconditioned Conjugate Gradient solver
pub struct PreconditionedCG;

impl PreconditionedCG {
    /// Solve Ax = b using Preconditioned Conjugate Gradient
    pub fn solve(
        a: &Array2<f64>,
        b: &Array1<f64>,
        max_iter: usize,
        tol: f64,
        preconditioner: &PreconditionerType,
    ) -> Result<Array1<f64>> {
        let n = a.nrows();
        let mut x = Array1::zeros(n);
        let mut r = b - &a.dot(&x);

        // Setup preconditioner
        let precond_matrix = PreconditionerSetup::setup_preconditioner(a, preconditioner)?;
        let mut z = PreconditionerSetup::apply_preconditioner(&precond_matrix, &r, preconditioner)?;
        let mut p = z.clone();
        let mut rsold = r.dot(&z);

        for _iter in 0..max_iter {
            let ap = a.dot(&p);
            let alpha = rsold / p.dot(&ap);

            x = &x + alpha * &p;
            r = &r - alpha * &ap;

            // Check convergence
            let rnorm = r.mapv(|x| x * x).sum().sqrt();
            if rnorm < tol {
                break;
            }

            z = PreconditionerSetup::apply_preconditioner(&precond_matrix, &r, preconditioner)?;
            let rsnew = r.dot(&z);
            let beta = rsnew / rsold;

            p = &z + beta * &p;
            rsold = rsnew;
        }

        Ok(x)
    }
}

/// Preconditioner setup and application
pub struct PreconditionerSetup;

impl PreconditionerSetup {
    /// Setup preconditioner matrix
    pub fn setup_preconditioner(
        a: &Array2<f64>,
        preconditioner: &PreconditionerType,
    ) -> Result<Array2<f64>> {
        match preconditioner {
            PreconditionerType::None => Ok(Array2::eye(a.nrows())),

            PreconditionerType::Diagonal => {
                // Diagonal preconditioner M = diag(A)
                let diag_inv = a
                    .diag()
                    .mapv(|x| if x.abs() > 1e-12 { 1.0 / x } else { 1.0 });
                Ok(Array2::from_diag(&diag_inv))
            }

            PreconditionerType::IncompleteCholesky { fill_factor: _ } => {
                // Simplified incomplete Cholesky (just return diagonal for now)
                let diag_inv = a
                    .diag()
                    .mapv(|x| if x > 1e-12 { 1.0 / x.sqrt() } else { 1.0 });
                Ok(Array2::from_diag(&diag_inv))
            }

            PreconditionerType::SSOR { omega } => {
                // SSOR preconditioner setup
                let n = a.nrows();
                let mut d = Array2::zeros((n, n));
                let mut l = Array2::zeros((n, n));

                // Extract diagonal and lower triangular parts
                for i in 0..n {
                    d[(i, i)] = a[(i, i)];
                    for j in 0..i {
                        l[(i, j)] = a[(i, j)];
                    }
                }

                // SSOR matrix: M = (D + omega*L) * D^(-1) * (D + omega*L)^T
                let d_inv =
                    Array2::from_diag(
                        &d.diag()
                            .mapv(|x| if x.abs() > 1e-12 { 1.0 / x } else { 1.0 }),
                    );
                let dl = &d + *omega * &l;
                Ok(dl.dot(&d_inv).dot(&dl.t()))
            }
        }
    }

    /// Apply preconditioner to vector
    pub fn apply_preconditioner(
        precond: &Array2<f64>,
        vector: &Array1<f64>,
        preconditioner: &PreconditionerType,
    ) -> Result<Array1<f64>> {
        match preconditioner {
            PreconditionerType::None => Ok(vector.clone()),
            PreconditionerType::Diagonal => Ok(&precond.diag().to_owned() * vector),
            _ => Ok(precond.dot(vector)),
        }
    }
}

/// Lanczos eigendecomposition method
pub struct LanczosMethod;

impl LanczosMethod {
    /// Perform Lanczos eigendecomposition
    pub fn eigendecomposition(
        matrix: &Array2<f64>,
        num_vectors: usize,
        tol: f64,
    ) -> Result<(Array1<f64>, Array2<f64>)> {
        let n = matrix.nrows();
        let m = num_vectors.min(n);

        // Initialize Lanczos vectors
        let mut q_matrix = Array2::zeros((n, m));
        let mut alpha_vec = Array1::zeros(m);
        let mut beta_vec = Array1::zeros(m);

        // Random starting vector
        let mut rng = thread_rng();
        let uniform = RandUniform::new(-1.0, 1.0).unwrap();
        let mut q_0 = Array1::zeros(n);
        for i in 0..n {
            q_0[i] = rng.sample(uniform);
        }
        #[allow(clippy::unnecessary_cast)]
        let q_0_norm = (q_0.mapv(|x| x * x).sum() as f64).sqrt();
        q_0 /= q_0_norm;
        q_matrix.column_mut(0).assign(&q_0);

        let mut beta = 0.0;
        let mut q_prev = Array1::zeros(n);

        for j in 0..m {
            let q_j = q_matrix.column(j).to_owned();
            let mut w: Array1<f64> = matrix.dot(&q_j) - beta * &q_prev;

            alpha_vec[j] = q_j.dot(&w);
            w = &w - alpha_vec[j] * &q_j;

            beta = w.mapv(|x| x * x).sum().sqrt();
            if j < m - 1 {
                beta_vec[j] = beta;
                if beta < tol {
                    break;
                }
                q_matrix.column_mut(j + 1).assign(&(&w / beta));
            }

            q_prev = q_j;
        }

        // Solve tridiagonal eigenvalue problem
        let (eigenvals, eigenvecs_tri) = TridiagonalEigenSolver::solve(&alpha_vec, &beta_vec)?;

        // Transform back to original space
        let eigenvecs = q_matrix.dot(&eigenvecs_tri);

        Ok((eigenvals, eigenvecs))
    }
}

/// Tridiagonal eigenvalue problem solver
pub struct TridiagonalEigenSolver;

impl TridiagonalEigenSolver {
    /// Solve tridiagonal eigenvalue problem using simplified QR algorithm
    pub fn solve(alpha: &Array1<f64>, beta: &Array1<f64>) -> Result<(Array1<f64>, Array2<f64>)> {
        let n = alpha.len();

        // Build tridiagonal matrix
        let mut tri_matrix = Array2::zeros((n, n));
        for i in 0..n {
            tri_matrix[(i, i)] = alpha[i];
            if i < n - 1 {
                tri_matrix[(i, i + 1)] = beta[i];
                tri_matrix[(i + 1, i)] = beta[i];
            }
        }

        // Use SVD for eigendecomposition (simplified approach)
        let (u, s, _vt) = tri_matrix
            .svd(true)
            .map_err(|e| SklearsError::NumericalError(format!("SVD failed: {:?}", e)))?;

        Ok((s, u))
    }
}

/// Iterative refinement for improved numerical accuracy
pub struct IterativeRefinement;

impl IterativeRefinement {
    /// Perform iterative refinement on a linear system solution
    pub fn refine_solution(
        a: &Array2<f64>,
        b: &Array1<f64>,
        x: &Array1<f64>,
        max_iter: usize,
        tol: f64,
    ) -> Result<Array1<f64>> {
        let mut x_refined = x.clone();

        for _iter in 0..max_iter {
            // Compute residual: r = b - A*x
            let residual = b - &a.dot(&x_refined);

            // Check convergence
            let residual_norm = residual.mapv(|x| x * x).sum().sqrt();
            if residual_norm < tol {
                break;
            }

            // Solve A*dx = r for correction
            let dx = KernelOps::invert_using_cholesky(a)?.dot(&residual);

            // Update solution
            x_refined = &x_refined + &dx;
        }

        Ok(x_refined)
    }
}

/// Specialized solvers for specific matrix structures
pub struct SpecializedSolvers;

impl SpecializedSolvers {
    /// Solve system with Kronecker product structure
    pub fn solve_kronecker(
        a1: &Array2<f64>,
        a2: &Array2<f64>,
        b: &Array2<f64>,
    ) -> Result<Array2<f64>> {
        // For system (A2 ⊗ A1) vec(X) = vec(B)
        // Solution is X = A1^(-1) * B * A2^(-T)

        let a1_inv = KernelOps::invert_using_cholesky(a1)?;
        let a2_inv = KernelOps::invert_using_cholesky(a2)?;

        let x = a1_inv.dot(b).dot(&a2_inv.t());

        Ok(x)
    }

    /// Solve system with block diagonal structure
    pub fn solve_block_diagonal(
        blocks: &[Array2<f64>],
        rhs_blocks: &[Array1<f64>],
    ) -> Result<Array1<f64>> {
        if blocks.len() != rhs_blocks.len() {
            return Err(SklearsError::InvalidInput(
                "Number of blocks must match RHS blocks".to_string(),
            ));
        }

        let mut solution_blocks = Vec::new();

        for (block, rhs_block) in blocks.iter().zip(rhs_blocks.iter()) {
            let block_inv = KernelOps::invert_using_cholesky(block)?;
            let block_solution = block_inv.dot(rhs_block);
            solution_blocks.push(block_solution);
        }

        // Concatenate solutions
        let total_size: usize = solution_blocks.iter().map(|b| b.len()).sum();
        let mut solution = Array1::zeros(total_size);
        let mut offset = 0;

        for block_solution in solution_blocks {
            let block_size = block_solution.len();
            solution
                .slice_mut(s![offset..offset + block_size])
                .assign(&block_solution);
            offset += block_size;
        }

        Ok(solution)
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;

    use approx::assert_abs_diff_eq;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_direct_prediction() {
        let k_star_m = array![[0.5, 0.3], [0.7, 0.2]];
        let alpha = array![1.0, 2.0];

        let result = ScalableInference::predict_direct(&k_star_m, &alpha).unwrap();
        let expected = array![1.1, 1.1]; // 0.5*1.0 + 0.3*2.0, 0.7*1.0 + 0.2*2.0

        for (a, b) in result.iter().zip(expected.iter()) {
            assert_abs_diff_eq!(*a, *b, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_diagonal_preconditioner() {
        let matrix = array![[4.0, 1.0], [1.0, 3.0]];
        let precond =
            PreconditionerSetup::setup_preconditioner(&matrix, &PreconditionerType::Diagonal)
                .unwrap();

        // Should be diag([1/4, 1/3])
        assert_abs_diff_eq!(precond[(0, 0)], 0.25, epsilon = 1e-10);
        assert_abs_diff_eq!(precond[(1, 1)], 1.0 / 3.0, epsilon = 1e-10);
        assert_abs_diff_eq!(precond[(0, 1)], 0.0, epsilon = 1e-10);
    }

    #[test]
    fn test_pcg_solver() {
        let a = array![[4.0, 1.0], [1.0, 3.0]];
        let b = array![1.0, 2.0];

        let solution =
            PreconditionedCG::solve(&a, &b, 100, 1e-10, &PreconditionerType::Diagonal).unwrap();

        // Verify A*x = b
        let residual = &b - &a.dot(&solution);
        let residual_norm = residual.mapv(|x| x * x).sum().sqrt();
        assert!(residual_norm < 1e-8);
    }

    #[test]
    fn test_lanczos_eigendecomposition() {
        let matrix = array![[3.0, 1.0], [1.0, 2.0]];

        let (eigenvals, eigenvecs) = LanczosMethod::eigendecomposition(&matrix, 2, 1e-10).unwrap();

        assert_eq!(eigenvals.len(), 2);
        assert_eq!(eigenvecs.shape(), &[2, 2]);

        // Eigenvalues should be positive for positive definite matrix
        assert!(eigenvals.iter().all(|&x| x > 0.0));
    }

    #[test]
    fn test_iterative_refinement() {
        let a = array![[2.0, 1.0], [1.0, 2.0]];
        let b = array![3.0, 3.0];
        let x_initial = array![1.0, 1.0]; // Exact solution

        let x_refined =
            IterativeRefinement::refine_solution(&a, &b, &x_initial, 10, 1e-12).unwrap();

        // Solution should remain close to initial (which is exact)
        for (a, b) in x_refined.iter().zip(x_initial.iter()) {
            assert_abs_diff_eq!(*a, *b, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_block_diagonal_solver() {
        let block1 = array![[2.0, 0.0], [0.0, 3.0]];
        let block2 = array![[1.0]];
        let blocks = vec![block1, block2];

        let rhs1 = array![4.0, 6.0];
        let rhs2 = array![2.0];
        let rhs_blocks = vec![rhs1, rhs2];

        let solution = SpecializedSolvers::solve_block_diagonal(&blocks, &rhs_blocks).unwrap();

        // Expected: [2.0, 2.0, 2.0] (4/2, 6/3, 2/1)
        let expected = array![2.0, 2.0, 2.0];
        assert_eq!(solution.len(), expected.len());

        for (a, b) in solution.iter().zip(expected.iter()) {
            assert_abs_diff_eq!(*a, *b, epsilon = 1e-5);
        }
    }
}