blr-active 0.1.0

//! Algorithm 1: Posterior Variance Computation
//!
//! Computes σ²_* = 1/β + φ_* Σ φ_*^T for arbitrary test points, where:
//! - β is the noise precision (1/σ²_noise)
//! - Σ is the posterior covariance matrix (D×D, row-major)
//! - φ_* is the feature vector of a test point
//!
//! Numerical stability: works directly on the already-computed posterior Σ
//! (no matrix inversion required here — inversion happens once during BLR fit).
//! Negative variances due to floating-point error are clamped to 1/β.

/// Compute posterior variance for each of `n_test` test points.
///
/// # Arguments
/// - `beta`: noise precision β
/// - `sigma_cov`: D×D posterior covariance matrix, row-major flat slice
/// - `d`: feature dimension D
/// - `phi_test`: N_test × D feature matrix, row-major flat slice
/// - `n_test`: number of test points
///
/// # Returns
/// Vec of length `n_test` with σ²_* for each point.
///
/// # Performance
/// O(n_test · D²) — for 100 points × 8 features ≈ 6400 ops, well under 1 ms.
pub fn posterior_variance(
    beta: f64,
    sigma_cov: &[f64],
    d: usize,
    phi_test: &[f64],
    n_test: usize,
) -> Vec<f64> {
    debug_assert_eq!(sigma_cov.len(), d * d, "sigma_cov must be D×D row-major");
    debug_assert_eq!(
        phi_test.len(),
        n_test * d,
        "phi_test must be N_test×D row-major"
    );

    let noise_var = 1.0 / beta.max(1e-10); // guard against β=0
    let mut variances = vec![noise_var; n_test];

    for i in 0..n_test {
        let phi_i = &phi_test[i * d..(i + 1) * d];

        // Compute Σ φ_i (D-vector): sigma_cov[row][col] * phi_i[col]
        let mut sigma_phi = vec![0.0_f64; d];
        for row in 0..d {
            let mut acc = 0.0_f64;
            for col in 0..d {
                acc += sigma_cov[row * d + col] * phi_i[col];
            }
            sigma_phi[row] = acc;
        }

        // epistemic_var = φ_i^T (Σ φ_i)
        let epistemic: f64 = phi_i.iter().zip(sigma_phi.iter()).map(|(a, b)| a * b).sum();

        // Clamp to zero: floating-point errors can produce tiny negatives
        // NUMERICAL GUARD: if epistemic < 0, keep only noise variance
        variances[i] = noise_var + epistemic.max(0.0);
    }

    variances
}

/// Compute posterior standard deviations for each of `n_test` test points.
///
/// Convenience wrapper around [`posterior_variance`] that takes the square root.
pub fn posterior_std(
    beta: f64,
    sigma_cov: &[f64],
    d: usize,
    phi_test: &[f64],
    n_test: usize,
) -> Vec<f64> {
    posterior_variance(beta, sigma_cov, d, phi_test, n_test)
        .into_iter()
        .map(f64::sqrt)
        .collect()
}

/// Compute posterior std over a uniform 1-D input grid.
///
/// For each grid point, calls `feature_fn(x)` to obtain a D-dimensional
/// feature vector, then evaluates Algorithm 1.
///
/// # Arguments
/// - `beta`: noise precision β
/// - `sigma_cov`: D×D posterior covariance, row-major flat slice
/// - `d`: feature dimension D
/// - `input_min`, `input_max`: inclusive range for the grid
/// - `resolution`: number of grid points (>= 2)
/// - `feature_fn`: maps a scalar input value to a `Vec<f64>` of length D
///
/// # Returns
/// `(grid_points, std_devs)` — both of length `resolution`.
///
/// # Panics
/// Does not panic; silently clamps degenerate feature vectors if D mismatches.
pub fn posterior_std_grid(
    beta: f64,
    sigma_cov: &[f64],
    d: usize,
    input_min: f64,
    input_max: f64,
    resolution: usize,
    feature_fn: &dyn Fn(f64) -> Vec<f64>,
) -> (Vec<f64>, Vec<f64>) {
    let resolution = resolution.max(2);
    let step = (input_max - input_min) / (resolution - 1) as f64;

    let grid: Vec<f64> = (0..resolution)
        .map(|k| input_min + k as f64 * step)
        .collect();

    // Build N×D feature matrix row-major
    let mut phi_grid = Vec::with_capacity(resolution * d);
    for &x in &grid {
        let feats = feature_fn(x);
        // Pad or truncate to D if feature_fn returns wrong dimensionality
        let actual_d = feats.len().min(d);
        phi_grid.extend_from_slice(&feats[..actual_d]);
        if actual_d < d {
            phi_grid.extend(std::iter::repeat_n(0.0, d - actual_d));
        }
    }

    let stds = posterior_std(beta, sigma_cov, d, &phi_grid, resolution);
    (grid, stds)
}

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

    /// Identity covariance Σ=I, β=1 → var at φ=(1,0,...) should be 1 + 1 = 2
    #[test]
    fn test_variance_identity_cov() {
        let d = 2usize;
        let sigma_cov = vec![1.0, 0.0, 0.0, 1.0]; // I₂
        let beta = 1.0;
        // φ = [1, 0]
        let phi_test = vec![1.0, 0.0];
        let var = posterior_variance(beta, &sigma_cov, d, &phi_test, 1);
        assert!((var[0] - 2.0).abs() < 1e-10, "expected 2.0, got {}", var[0]);
    }

    /// All variances should be finite and non-negative everywhere on a grid
    #[test]
    fn test_variance_grid_nonnegative() {
        let d = 3usize;
        // Diagonal covariance with varying values
        let sigma_cov = vec![0.1, 0.0, 0.0, 0.0, 0.2, 0.0, 0.0, 0.0, 0.05];
        let beta = 2.0;
        let feature_fn = |x: f64| vec![x, x * x, x.sin()];
        let (_, stds) = posterior_std_grid(beta, &sigma_cov, d, -5.0, 5.0, 50, &feature_fn);
        for s in &stds {
            assert!(
                s.is_finite() && *s >= 0.0,
                "non-finite or negative std: {s}"
            );
        }
    }

    /// Same input should always produce same output (determinism)
    #[test]
    fn test_variance_deterministic() {
        let d = 2usize;
        let sigma_cov = vec![0.5, 0.1, 0.1, 0.3];
        let beta = 1.5;
        let phi_test = vec![1.0, -0.5, 0.0, 1.0];
        let var1 = posterior_variance(beta, &sigma_cov, d, &phi_test, 2);
        let var2 = posterior_variance(beta, &sigma_cov, d, &phi_test, 2);
        for (a, b) in var1.iter().zip(var2.iter()) {
            assert_eq!(a, b, "result not deterministic");
        }
    }

    /// Near-singular (but PSD) covariance should not crash
    #[test]
    fn test_variance_near_singular() {
        let d = 3usize;
        // Nearly degenerate Σ — tiny eigenvalue on last dimension
        let eps = 1e-9;
        let sigma_cov = vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, eps];
        let beta = 1.0;
        let phi_test = vec![0.0, 0.0, 100.0]; // stress the near-zero dim
        let var = posterior_variance(beta, &sigma_cov, d, &phi_test, 1);
        assert!(var[0].is_finite(), "variance must be finite: {}", var[0]);
        assert!(var[0] >= 0.0, "variance must be non-negative: {}", var[0]);
    }

    /// β=0 (edge case) — should not divide by zero; noise_var clamped to 1/1e-10
    #[test]
    fn test_variance_zero_beta() {
        let d = 2usize;
        let sigma_cov = vec![1.0, 0.0, 0.0, 1.0];
        let var = posterior_variance(0.0, &sigma_cov, d, &[0.5, 0.5], 1);
        assert!(var[0].is_finite(), "variance must be finite even with β=0");
    }
}