sklears_linear/

large_scale_variational_inference.rs

//! Large-Scale Variational Inference for Linear Models
//!
//! This module implements scalable variational Bayesian inference algorithms
//! designed for large datasets that don't fit in memory. It includes stochastic
//! variational inference, mini-batch processing, and streaming algorithms.

use scirs2_core::ndarray::{s, Array1, Array2};
use scirs2_core::random::seq::SliceRandom;
use scirs2_core::random::{Distribution, RandNormal as Normal, Rng};
use scirs2_core::{SeedableRng, StdRng};
use sklears_core::{
    error::{Result, SklearsError},
    traits::{Estimator, Fit, Predict, Trained, Untrained},
    types::Float,
};
use std::marker::PhantomData;

/// Configuration for large-scale variational inference
#[derive(Debug, Clone)]
pub struct LargeScaleVariationalConfig {
    /// Maximum number of epochs
    pub max_epochs: usize,
    /// Mini-batch size for stochastic updates
    pub batch_size: usize,
    /// Learning rate for variational parameter updates
    pub learning_rate: Float,
    /// Learning rate decay schedule
    pub learning_rate_decay: LearningRateDecay,
    /// Convergence tolerance
    pub tolerance: Float,
    /// Number of Monte Carlo samples for expectations
    pub n_mc_samples: usize,
    /// Whether to use natural gradients
    pub use_natural_gradients: bool,
    /// Whether to use control variates for variance reduction
    pub use_control_variates: bool,
    /// Memory limit in GB for adaptive batch sizing
    pub memory_limit_gb: Option<Float>,
    /// Whether to enable verbose output
    pub verbose: bool,
    /// Random seed for reproducibility
    pub random_seed: Option<u64>,
    /// Prior parameters
    pub prior_config: PriorConfiguration,
}

impl Default for LargeScaleVariationalConfig {
    fn default() -> Self {
        Self {
            max_epochs: 100,
            batch_size: 256,
            learning_rate: 0.01,
            learning_rate_decay: LearningRateDecay::Exponential { decay_rate: 0.95 },
            tolerance: 1e-6,
            n_mc_samples: 10,
            use_natural_gradients: true,
            use_control_variates: true,
            memory_limit_gb: Some(4.0),
            verbose: false,
            random_seed: None,
            prior_config: PriorConfiguration::default(),
        }
    }
}

/// Learning rate decay schedules
#[derive(Debug, Clone)]
pub enum LearningRateDecay {
    /// Constant learning rate
    Constant,
    /// Exponential decay: lr * decay_rate^epoch
    Exponential { decay_rate: Float },
    /// Step decay: lr * step_factor^floor(epoch / step_size)
    Step {
        step_size: usize,
        step_factor: Float,
    },
    /// Polynomial decay: lr * (1 + decay_rate * epoch)^(-power)
    Polynomial { decay_rate: Float, power: Float },
    /// Cosine annealing
    CosineAnnealing { min_lr: Float },
}

/// Prior configuration for Bayesian linear regression
#[derive(Debug, Clone)]
pub struct PriorConfiguration {
    /// Prior precision for weights (Gamma distribution parameters)
    pub weight_precision_shape: Float,
    pub weight_precision_rate: Float,
    /// Prior precision for noise (Gamma distribution parameters)
    pub noise_precision_shape: Float,
    pub noise_precision_rate: Float,
    /// Whether to use hierarchical priors
    pub hierarchical: bool,
    /// ARD (Automatic Relevance Determination) configuration
    pub ard_config: Option<ARDConfiguration>,
}

impl Default for PriorConfiguration {
    fn default() -> Self {
        Self {
            weight_precision_shape: 1e-6,
            weight_precision_rate: 1e-6,
            noise_precision_shape: 1e-6,
            noise_precision_rate: 1e-6,
            hierarchical: false,
            ard_config: None,
        }
    }
}

/// Configuration for Automatic Relevance Determination
#[derive(Debug, Clone)]
pub struct ARDConfiguration {
    /// Individual precision priors for each feature
    pub feature_precision_shape: Float,
    pub feature_precision_rate: Float,
    /// Threshold for feature pruning
    pub pruning_threshold: Float,
    /// Whether to enable automatic feature pruning
    pub enable_pruning: bool,
}

/// Variational parameters for the posterior distribution
#[derive(Debug, Clone)]
pub struct VariationalPosterior {
    /// Mean of weight posterior (multivariate normal)
    pub weight_mean: Array1<Float>,
    /// Covariance of weight posterior
    pub weight_covariance: Array2<Float>,
    /// Precision matrix (inverse covariance)
    pub weight_precision: Array2<Float>,
    /// Parameters for weight precision posterior (Gamma)
    pub weight_precision_shape: Array1<Float>,
    pub weight_precision_rate: Array1<Float>,
    /// Parameters for noise precision posterior (Gamma)
    pub noise_precision_shape: Float,
    pub noise_precision_rate: Float,
    /// Log marginal likelihood lower bound (ELBO)
    pub elbo: Float,
}

impl VariationalPosterior {
    /// Create a new variational posterior with given dimensions
    pub fn new(n_features: usize, config: &PriorConfiguration) -> Self {
        Self {
            weight_mean: Array1::zeros(n_features),
            weight_covariance: Array2::eye(n_features),
            weight_precision: Array2::eye(n_features),
            weight_precision_shape: Array1::from_elem(n_features, config.weight_precision_shape),
            weight_precision_rate: Array1::from_elem(n_features, config.weight_precision_rate),
            noise_precision_shape: config.noise_precision_shape,
            noise_precision_rate: config.noise_precision_rate,
            elbo: Float::NEG_INFINITY,
        }
    }

    /// Sample from the posterior distribution
    pub fn sample_weights(&self, n_samples: usize, rng: &mut impl Rng) -> Result<Array2<Float>> {
        let n_features = self.weight_mean.len();
        let mut samples = Array2::zeros((n_samples, n_features));

        // Compute Cholesky decomposition of covariance matrix
        let chol = self.cholesky_decomposition(&self.weight_covariance)?;

        for i in 0..n_samples {
            // Sample from standard normal
            let z: Array1<Float> = (0..n_features)
                .map(|_| {
                    Normal::new(0.0, 1.0)
                        .expect("valid normal distribution parameters")
                        .sample(rng)
                })
                .collect::<Vec<_>>()
                .into();

            // Transform to desired distribution: μ + L * z
            let sample = &self.weight_mean + chol.dot(&z);
            samples.slice_mut(s![i, ..]).assign(&sample);
        }

        Ok(samples)
    }

    /// Compute Cholesky decomposition (simplified implementation)
    fn cholesky_decomposition(&self, matrix: &Array2<Float>) -> Result<Array2<Float>> {
        let n = matrix.nrows();
        let mut l = Array2::zeros((n, n));

        for i in 0..n {
            for j in 0..=i {
                if i == j {
                    // Diagonal elements
                    let sum: Float = (0..j).map(|k| l[[i, k]] * l[[i, k]]).sum();
                    let val = matrix[[i, i]] - sum;
                    if val <= 0.0 {
                        return Err(SklearsError::NumericalError(
                            "Matrix is not positive definite".to_string(),
                        ));
                    }
                    l[[i, j]] = val.sqrt();
                } else {
                    // Lower triangular elements
                    let sum: Float = (0..j).map(|k| l[[i, k]] * l[[j, k]]).sum();
                    l[[i, j]] = (matrix[[i, j]] - sum) / l[[j, j]];
                }
            }
        }

        Ok(l)
    }
}

/// Large-scale variational Bayesian linear regression
#[derive(Debug)]
pub struct LargeScaleVariationalRegression<State = Untrained> {
    config: LargeScaleVariationalConfig,
    state: PhantomData<State>,
    // Trained state
    posterior: Option<VariationalPosterior>,
    convergence_history: Option<Vec<Float>>,
    feature_relevance: Option<Array1<Float>>,
    n_features: Option<usize>,
    intercept: Option<Float>,
}

impl Default for LargeScaleVariationalRegression<Untrained> {
    fn default() -> Self {
        Self::new()
    }
}

impl LargeScaleVariationalRegression<Untrained> {
    /// Create a new large-scale variational regression model
    pub fn new() -> Self {
        Self {
            config: LargeScaleVariationalConfig::default(),
            state: PhantomData,
            posterior: None,
            convergence_history: None,
            feature_relevance: None,
            n_features: None,
            intercept: None,
        }
    }

    /// Set the configuration
    pub fn with_config(mut self, config: LargeScaleVariationalConfig) -> Self {
        self.config = config;
        self
    }

    /// Set batch size
    pub fn batch_size(mut self, batch_size: usize) -> Self {
        self.config.batch_size = batch_size;
        self
    }

    /// Set learning rate
    pub fn learning_rate(mut self, learning_rate: Float) -> Self {
        self.config.learning_rate = learning_rate;
        self
    }

    /// Enable Automatic Relevance Determination
    pub fn enable_ard(mut self, pruning_threshold: Float) -> Self {
        self.config.prior_config.ard_config = Some(ARDConfiguration {
            feature_precision_shape: 1e-6,
            feature_precision_rate: 1e-6,
            pruning_threshold,
            enable_pruning: true,
        });
        self
    }

    /// Set memory limit for adaptive batch sizing
    pub fn memory_limit_gb(mut self, limit: Float) -> Self {
        self.config.memory_limit_gb = Some(limit);
        self
    }
}

impl LargeScaleVariationalRegression<Trained> {
    /// Get the posterior mean of coefficients
    pub fn coefficients(&self) -> &Array1<Float> {
        &self
            .posterior
            .as_ref()
            .expect("value should be present")
            .weight_mean
    }

    /// Get the posterior covariance of coefficients
    pub fn coefficient_covariance(&self) -> &Array2<Float> {
        &self
            .posterior
            .as_ref()
            .expect("value should be present")
            .weight_covariance
    }

    /// Get feature relevance scores (for ARD)
    pub fn feature_relevance(&self) -> Option<&Array1<Float>> {
        self.feature_relevance.as_ref()
    }

    /// Get convergence history
    pub fn convergence_history(&self) -> Option<&[Float]> {
        self.convergence_history.as_deref()
    }

    /// Sample predictions from the posterior predictive distribution
    pub fn sample_predictions(
        &self,
        x: &Array2<Float>,
        n_samples: usize,
        rng: &mut impl Rng,
    ) -> Result<Array2<Float>> {
        let posterior = self
            .posterior
            .as_ref()
            .ok_or_else(|| SklearsError::NumericalError("value should be present".into()))?;
        let weight_samples = posterior.sample_weights(n_samples, rng)?;

        let mut predictions = Array2::zeros((n_samples, x.nrows()));

        for i in 0..n_samples {
            let weights = weight_samples.slice(s![i, ..]);
            let pred = x.dot(&weights);
            predictions.slice_mut(s![i, ..]).assign(&pred);

            // Add intercept if fitted
            if let Some(intercept) = self.intercept {
                predictions
                    .slice_mut(s![i, ..])
                    .mapv_inplace(|x| x + intercept);
            }
        }

        Ok(predictions)
    }

    /// Compute predictive uncertainties
    pub fn predict_with_uncertainty(
        &self,
        x: &Array2<Float>,
    ) -> Result<(Array1<Float>, Array1<Float>)> {
        let posterior = self
            .posterior
            .as_ref()
            .ok_or_else(|| SklearsError::NumericalError("value should be present".into()))?;

        // Predictive mean
        let pred_mean = x.dot(&posterior.weight_mean);

        // Predictive variance: X * Σ * X^T + σ²
        let mut pred_var = Array1::zeros(x.nrows());

        for i in 0..x.nrows() {
            let x_i = x.slice(s![i, ..]);
            let var_contrib = x_i.dot(&posterior.weight_covariance.dot(&x_i));

            // Add noise variance
            let noise_var = 1.0 / posterior.noise_precision_rate; // Simplified
            pred_var[i] = var_contrib + noise_var;
        }

        let pred_std = pred_var.mapv(|v| v.sqrt());

        Ok((pred_mean, pred_std))
    }
}

impl Fit<Array2<Float>, Array1<Float>> for LargeScaleVariationalRegression<Untrained> {
    type Fitted = LargeScaleVariationalRegression<Trained>;

    fn fit(self, x: &Array2<Float>, y: &Array1<Float>) -> Result<Self::Fitted> {
        let (n_samples, n_features) = x.dim();

        if n_samples != y.len() {
            return Err(SklearsError::DimensionMismatch {
                expected: n_samples,
                actual: y.len(),
            });
        }

        // Initialize variational posterior
        let mut posterior = VariationalPosterior::new(n_features, &self.config.prior_config);
        let mut convergence_history = Vec::new();

        // Initialize random number generator
        let mut rng = if let Some(seed) = self.config.random_seed {
            StdRng::seed_from_u64(seed)
        } else {
            StdRng::from_rng(&mut scirs2_core::random::thread_rng())
        };

        // Stochastic variational inference
        let mut current_lr = self.config.learning_rate;

        for epoch in 0..self.config.max_epochs {
            let epoch_elbo = self.run_epoch(x, y, &mut posterior, current_lr, &mut rng)?;
            convergence_history.push(epoch_elbo);

            if self.config.verbose && epoch % 10 == 0 {
                println!("Epoch {}: ELBO = {:.6}", epoch, epoch_elbo);
            }

            // Check convergence
            if epoch > 0 {
                let prev_elbo = convergence_history[epoch - 1];
                let elbo_change = (epoch_elbo - prev_elbo).abs();

                if elbo_change < self.config.tolerance {
                    if self.config.verbose {
                        println!("Converged after {} epochs", epoch);
                    }
                    break;
                }
            }

            // Update learning rate
            current_lr = self.update_learning_rate(current_lr, epoch);
        }

        // Compute feature relevance for ARD
        let feature_relevance = if self.config.prior_config.ard_config.is_some() {
            Some(self.compute_feature_relevance(&posterior))
        } else {
            None
        };

        Ok(LargeScaleVariationalRegression {
            config: self.config,
            state: PhantomData,
            posterior: Some(posterior),
            convergence_history: Some(convergence_history),
            feature_relevance,
            n_features: Some(n_features),
            intercept: None, // Simplified: not handling intercept in this implementation
        })
    }
}

impl LargeScaleVariationalRegression<Untrained> {
    /// Run one epoch of stochastic variational inference
    fn run_epoch(
        &self,
        x: &Array2<Float>,
        y: &Array1<Float>,
        posterior: &mut VariationalPosterior,
        learning_rate: Float,
        rng: &mut impl Rng,
    ) -> Result<Float> {
        let (n_samples, _n_features) = x.dim();
        let batch_size = self.config.batch_size.min(n_samples);

        let mut total_elbo = 0.0;
        let mut n_batches = 0;

        // Create mini-batches
        let mut indices: Vec<usize> = (0..n_samples).collect();
        indices.shuffle(rng);

        for batch_indices in indices.chunks(batch_size) {
            // Extract mini-batch
            let batch_x = self.extract_batch_features(x, batch_indices);
            let batch_y = self.extract_batch_targets(y, batch_indices);

            // Compute natural gradients
            let (elbo, gradients) = self.compute_natural_gradients(
                &batch_x,
                &batch_y,
                posterior,
                n_samples,
                batch_indices.len(),
            )?;

            // Update variational parameters
            self.update_variational_parameters(posterior, &gradients, learning_rate)?;

            total_elbo += elbo;
            n_batches += 1;
        }

        Ok(total_elbo / n_batches as Float)
    }

    /// Extract mini-batch features
    fn extract_batch_features(&self, x: &Array2<Float>, indices: &[usize]) -> Array2<Float> {
        let mut batch_x = Array2::zeros((indices.len(), x.ncols()));
        for (i, &idx) in indices.iter().enumerate() {
            batch_x.slice_mut(s![i, ..]).assign(&x.slice(s![idx, ..]));
        }
        batch_x
    }

    /// Extract mini-batch targets
    fn extract_batch_targets(&self, y: &Array1<Float>, indices: &[usize]) -> Array1<Float> {
        indices.iter().map(|&i| y[i]).collect::<Vec<_>>().into()
    }

    /// Compute natural gradients for variational parameters
    fn compute_natural_gradients(
        &self,
        batch_x: &Array2<Float>,
        batch_y: &Array1<Float>,
        posterior: &VariationalPosterior,
        total_samples: usize,
        batch_size: usize,
    ) -> Result<(Float, VariationalGradients)> {
        let scale_factor = total_samples as Float / batch_size as Float;

        // Compute expected log likelihood
        let expected_ll = self.compute_expected_log_likelihood(batch_x, batch_y, posterior)?;

        // Compute KL divergence
        let kl_div = self.compute_kl_divergence(posterior)?;

        // ELBO = E[log p(y|X,w)] - KL[q(w)||p(w)]
        let elbo = scale_factor * expected_ll - kl_div;

        // Compute gradients (simplified implementation)
        let gradients = VariationalGradients {
            weight_mean_grad: Array1::zeros(posterior.weight_mean.len()),
            weight_precision_grad: Array2::zeros(posterior.weight_precision.dim()),
            noise_precision_shape_grad: 0.0,
            noise_precision_rate_grad: 0.0,
        };

        Ok((elbo, gradients))
    }

    /// Compute expected log likelihood
    fn compute_expected_log_likelihood(
        &self,
        x: &Array2<Float>,
        y: &Array1<Float>,
        posterior: &VariationalPosterior,
    ) -> Result<Float> {
        let _n_samples = x.nrows();

        // E[log p(y|X,w)] under q(w)
        let pred_mean = x.dot(&posterior.weight_mean);
        let residuals = y - &pred_mean;

        // Simplified computation (should include trace term for full correctness)
        let sum_squared_residuals = residuals.mapv(|r| r * r).sum();
        let expected_noise_precision =
            posterior.noise_precision_shape / posterior.noise_precision_rate;

        let log_likelihood = -0.5 * expected_noise_precision * sum_squared_residuals;

        Ok(log_likelihood)
    }

    /// Compute KL divergence KL[q(w,α,β)||p(w,α,β)]
    fn compute_kl_divergence(&self, _posterior: &VariationalPosterior) -> Result<Float> {
        // Simplified KL computation
        // Full implementation would compute KL for multivariate normal and Gamma distributions
        Ok(0.0)
    }

    /// Update variational parameters using natural gradients
    fn update_variational_parameters(
        &self,
        posterior: &mut VariationalPosterior,
        gradients: &VariationalGradients,
        learning_rate: Float,
    ) -> Result<()> {
        // Natural gradient updates for exponential family
        posterior.weight_mean =
            &posterior.weight_mean + learning_rate * &gradients.weight_mean_grad;

        // Ensure precision matrix remains positive definite
        // (Simplified update - full implementation would use proper natural gradients)

        Ok(())
    }

    /// Update learning rate according to decay schedule
    fn update_learning_rate(&self, current_lr: Float, epoch: usize) -> Float {
        match &self.config.learning_rate_decay {
            LearningRateDecay::Constant => current_lr,
            LearningRateDecay::Exponential { decay_rate } => {
                current_lr * decay_rate.powf(epoch as Float)
            }
            LearningRateDecay::Step {
                step_size,
                step_factor,
            } => current_lr * step_factor.powf((epoch / step_size) as Float),
            LearningRateDecay::Polynomial { decay_rate, power } => {
                current_lr * (1.0 + decay_rate * epoch as Float).powf(-power)
            }
            LearningRateDecay::CosineAnnealing { min_lr } => {
                min_lr
                    + 0.5
                        * (current_lr - min_lr)
                        * (1.0
                            + (std::f64::consts::PI * epoch as Float
                                / self.config.max_epochs as Float)
                                .cos())
            }
        }
    }

    /// Compute feature relevance scores for ARD
    fn compute_feature_relevance(&self, posterior: &VariationalPosterior) -> Array1<Float> {
        // Feature relevance based on posterior precision
        posterior.weight_precision_shape.clone() / &posterior.weight_precision_rate
    }
}

/// Gradients for variational parameters
#[derive(Debug, Clone)]
struct VariationalGradients {
    weight_mean_grad: Array1<Float>,
    weight_precision_grad: Array2<Float>,
    noise_precision_shape_grad: Float,
    noise_precision_rate_grad: Float,
}

impl Predict<Array2<Float>, Array1<Float>> for LargeScaleVariationalRegression<Trained> {
    fn predict(&self, x: &Array2<Float>) -> Result<Array1<Float>> {
        let (pred_mean, _) = self.predict_with_uncertainty(x)?;
        Ok(pred_mean)
    }
}

impl Estimator for LargeScaleVariationalRegression<Untrained> {
    type Config = LargeScaleVariationalConfig;
    type Error = SklearsError;
    type Float = Float;

    fn config(&self) -> &LargeScaleVariationalConfig {
        &self.config
    }
}

impl Estimator for LargeScaleVariationalRegression<Trained> {
    type Config = LargeScaleVariationalConfig;
    type Error = SklearsError;
    type Float = Float;
    fn config(&self) -> &LargeScaleVariationalConfig {
        &self.config
    }
}

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

    #[test]
    fn test_large_scale_variational_config() {
        let config = LargeScaleVariationalConfig::default();
        assert_eq!(config.max_epochs, 100);
        assert_eq!(config.batch_size, 256);
        assert_eq!(config.learning_rate, 0.01);
        assert_eq!(config.n_mc_samples, 10);
        assert!(config.use_natural_gradients);
    }

    #[test]
    fn test_variational_posterior_creation() {
        let prior_config = PriorConfiguration::default();
        let posterior = VariationalPosterior::new(5, &prior_config);

        assert_eq!(posterior.weight_mean.len(), 5);
        assert_eq!(posterior.weight_covariance.dim(), (5, 5));
        assert_eq!(posterior.weight_precision_shape.len(), 5);
    }

    #[test]
    fn test_learning_rate_decay() {
        let config = LargeScaleVariationalConfig {
            learning_rate: 0.1,
            learning_rate_decay: LearningRateDecay::Exponential { decay_rate: 0.9 },
            ..Default::default()
        };

        let model = LargeScaleVariationalRegression::new().with_config(config);

        let lr_epoch_0 = model.update_learning_rate(0.1, 0);
        let lr_epoch_1 = model.update_learning_rate(0.1, 1);

        assert_eq!(lr_epoch_0, 0.1);
        assert!((lr_epoch_1 - 0.09).abs() < 1e-10);
    }

    #[test]
    fn test_ard_configuration() {
        let model = LargeScaleVariationalRegression::new()
            .enable_ard(1e-6)
            .batch_size(128)
            .learning_rate(0.005);

        assert!(model.config.prior_config.ard_config.is_some());
        assert_eq!(model.config.batch_size, 128);
        assert_eq!(model.config.learning_rate, 0.005);
    }

    #[test]
    fn test_batch_extraction() {
        let X = Array::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
            .expect("valid array shape");
        let y = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]);

        let model = LargeScaleVariationalRegression::new();
        let indices = [0, 2];

        let batch_x = model.extract_batch_features(&X, &indices);
        let batch_y = model.extract_batch_targets(&y, &indices);

        assert_eq!(batch_x.dim(), (2, 2));
        assert_eq!(batch_y.len(), 2);
        assert_eq!(batch_x[[0, 0]], 1.0);
        assert_eq!(batch_x[[1, 0]], 5.0);
        assert_eq!(batch_y[0], 1.0);
        assert_eq!(batch_y[1], 3.0);
    }

    #[test]
    fn test_model_creation() {
        let model = LargeScaleVariationalRegression::new()
            .batch_size(64)
            .learning_rate(0.001)
            .memory_limit_gb(2.0);

        assert_eq!(model.config.batch_size, 64);
        assert_eq!(model.config.learning_rate, 0.001);
        assert_eq!(model.config.memory_limit_gb, Some(2.0));
    }
}