scirs2-stats 0.6.5

//! Real training and posterior-predictive inference for
//! [`super::BayesianNeuralNetwork`].
//!
//! `fit` trains a small **deep ensemble** (Lakshminarayanan et al., 2017):
//! each ensemble member is initialized by sampling its weights and biases
//! from the network's declared priors (via a real RNG, not a fixed
//! constant), independently bootstrap-resampled from the training data, and
//! then trained by ordinary gradient descent using an exact, hand-derived
//! backpropagation pass (`forward_with_cache` / `backward`) through whatever
//! mix of `ReLU`/`Sigmoid`/`Tanh`/`Swish`/`GELU` activations the network was
//! built with. Random initialization plus bootstrap resampling gives each
//! ensemble member a genuinely different local optimum, and the resulting
//! ensemble's empirical mean/variance is a real, well-established
//! (if approximate) way to quantify a neural network's epistemic uncertainty.
//!
//! `predict_with_uncertainty` then does real forward passes through that
//! trained ensemble (or, if `fit` has not been called, through independent
//! prior draws of the whole network -- a real prior-predictive Monte Carlo
//! estimate) and reports the empirical mean/variance across the draws.

use super::{ActivationType, AdvancedBayesianFloat, BayesianNeuralNetwork, DistributionType};
use crate::error::{StatsError, StatsResult};
use scirs2_core::ndarray::{Array1, Array2, ArrayView2, Axis};
use scirs2_core::random::{Rng, RngExt, SeedableRng};
use scirs2_core::validation::checkarray_finite;

/// Configuration for [`BayesianNeuralNetwork::fit`]'s deep-ensemble training.
#[derive(Debug, Clone)]
pub struct BnnTrainingConfig {
    /// Number of independently-initialized-and-trained ensemble members.
    pub n_ensemble: usize,
    /// Number of full-batch gradient descent epochs per ensemble member.
    pub epochs: usize,
    /// Gradient descent learning rate.
    pub learning_rate: f64,
    /// Whether to train each ensemble member on an independent bootstrap
    /// resample of the training data (in addition to independent random
    /// weight initialization).
    pub bootstrap: bool,
    /// Optional RNG seed for reproducibility.
    pub seed: Option<u64>,
}

impl Default for BnnTrainingConfig {
    fn default() -> Self {
        Self {
            n_ensemble: 16,
            epochs: 200,
            learning_rate: 0.05,
            bootstrap: true,
            seed: None,
        }
    }
}

/// Sample a real `N(mean, 1/precision)` draw using `rng` (replaces the old
/// hardcoded-`u1 = u2 = 0.5` "Box-Muller transform" that always produced the
/// same value regardless of `rng` state).
fn sample_normal<F: AdvancedBayesianFloat, R: Rng + ?Sized>(
    mean: F,
    precision: F,
    rng: &mut R,
) -> F {
    use scirs2_core::random::{Distribution, StandardNormal};
    let eps = F::from(1e-12).expect("1e-12 fits in any Float");
    let std_dev = F::one() / precision.max(eps).sqrt();
    let z64: f64 = StandardNormal.sample(rng);
    let z = F::from(z64).unwrap_or(F::zero());
    mean + std_dev * z
}

fn prior_precision<F: AdvancedBayesianFloat>(prior: &DistributionType<F>) -> F {
    match prior {
        DistributionType::Normal { precision, .. } => *precision,
        _ => F::one(),
    }
}

impl<F: AdvancedBayesianFloat> BayesianNeuralNetwork<F> {
    /// Derivative of the layer activation function with respect to its
    /// pre-activation input `z` (i.e. `d/dz apply_activation(z, activation)`).
    fn activation_derivative(&self, z: F, activation: ActivationType) -> F {
        match activation {
            ActivationType::ReLU => {
                if z > F::zero() {
                    F::one()
                } else {
                    F::zero()
                }
            }
            ActivationType::Sigmoid => {
                let s = F::one() / (F::one() + (-z).exp());
                s * (F::one() - s)
            }
            ActivationType::Tanh => {
                let t = z.tanh();
                F::one() - t * t
            }
            ActivationType::Swish => {
                let s = F::one() / (F::one() + (-z).exp());
                s + z * s * (F::one() - s)
            }
            ActivationType::GELU => {
                let sqrt_2_pi = F::from(0.7978845608).expect("sqrt(2/pi) fits in any Float");
                let coeff = F::from(0.044715).expect("0.044715 fits in any Float");
                let three = F::from(3.0).expect("3.0 fits in any Float");
                let half = F::from(0.5).expect("0.5 fits in any Float");
                let g = sqrt_2_pi * (z + coeff * z * z * z);
                let g_prime = sqrt_2_pi * (F::one() + three * coeff * z * z);
                let tanh_g = g.tanh();
                half * (F::one() + tanh_g) + half * z * (F::one() - tanh_g * tanh_g) * g_prime
            }
        }
    }

    /// Forward pass that also returns every layer's pre-activation (`zs`)
    /// and post-activation (`acts`, with `acts[0] == x`) values, needed by
    /// [`Self::backward`].
    fn forward_with_cache(
        &self,
        x: &ArrayView2<F>,
        weights: &[Array2<F>],
        biases: &[Array1<F>],
    ) -> StatsResult<(Vec<Array2<F>>, Vec<Array2<F>>)> {
        let mut acts = vec![x.to_owned()];
        let mut zs = Vec::with_capacity(self.activations.len());
        for (layer_idx, &activation_type) in self.activations.iter().enumerate() {
            let z = self.linear_transform(
                &acts[layer_idx].view(),
                &weights[layer_idx],
                &biases[layer_idx],
            )?;
            let a = z.mapv(|val| self.apply_activation(val, activation_type));
            zs.push(z);
            acts.push(a);
        }
        Ok((zs, acts))
    }

    /// Exact backpropagation of the mean-squared-error loss between the
    /// network's (post-activation) output and `y`, returning the gradient of
    /// each layer's weight matrix and bias vector.
    fn backward(
        &self,
        x: &ArrayView2<F>,
        y: &ArrayView2<F>,
        weights: &[Array2<F>],
        biases: &[Array1<F>],
    ) -> StatsResult<(Vec<Array2<F>>, Vec<Array1<F>>)> {
        let (zs, acts) = self.forward_with_cache(x, weights, biases)?;
        let n_layers = self.activations.len();
        let n_samples_ = x.nrows();
        let output_dim = *self.architecture.last().ok_or_else(|| {
            StatsError::InvalidArgument("architecture must be non-empty".to_string())
        })?;

        let scale = F::from(2.0).expect("2.0 fits in any Float")
            / F::from((n_samples_ * output_dim).max(1)).expect("count fits in any Float");

        let y_pred = &acts[n_layers];
        let mut delta_a = (y_pred - y).mapv(|v| v * scale);

        let mut grads_w: Vec<Array2<F>> = (0..n_layers).map(|_| Array2::zeros((0, 0))).collect();
        let mut grads_b: Vec<Array1<F>> = (0..n_layers).map(|_| Array1::zeros(0)).collect();

        for l in (0..n_layers).rev() {
            let z_l = &zs[l];
            let mut delta_z = Array2::<F>::zeros(delta_a.raw_dim());
            for (dz, (da, zv)) in delta_z.iter_mut().zip(delta_a.iter().zip(z_l.iter())) {
                *dz = *da * self.activation_derivative(*zv, self.activations[l]);
            }

            let a_l = &acts[l];
            grads_w[l] = a_l.t().dot(&delta_z);
            grads_b[l] = delta_z.sum_axis(Axis(0));

            if l > 0 {
                delta_a = delta_z.dot(&weights[l].t());
            }
        }

        Ok((grads_w, grads_b))
    }

    /// Train a deep ensemble of posterior weight/bias samples: each member is
    /// randomly initialized from the network's priors, optionally
    /// bootstrap-resampled, and trained to a local optimum of the mean
    /// squared error via exact backpropagation + gradient descent. Populates
    /// `self.weight_samples` / `self.bias_samples`, which
    /// [`Self::predict_with_uncertainty`] then draws from.
    pub fn fit(
        &mut self,
        x: &ArrayView2<F>,
        y: &ArrayView2<F>,
        config: &BnnTrainingConfig,
    ) -> StatsResult<()> {
        checkarray_finite(x, "x")?;
        checkarray_finite(y, "y")?;
        if x.nrows() != y.nrows() {
            return Err(StatsError::DimensionMismatch(
                "x and y must have the same number of rows".to_string(),
            ));
        }
        let output_dim = *self.architecture.last().ok_or_else(|| {
            StatsError::InvalidArgument("architecture must be non-empty".to_string())
        })?;
        if y.ncols() != output_dim {
            return Err(StatsError::DimensionMismatch(format!(
                "y has {} columns, expected {} to match the network's output layer",
                y.ncols(),
                output_dim
            )));
        }
        if x.ncols() != self.architecture[0] {
            return Err(StatsError::DimensionMismatch(format!(
                "x has {} columns, expected {} to match the network's input layer",
                x.ncols(),
                self.architecture[0]
            )));
        }
        if config.n_ensemble == 0 {
            return Err(StatsError::InvalidArgument(
                "n_ensemble must be at least 1".to_string(),
            ));
        }

        let mut rng = match config.seed {
            Some(seed) => scirs2_core::random::rngs::StdRng::seed_from_u64(seed),
            None => {
                scirs2_core::random::rngs::StdRng::from_rng(&mut scirs2_core::random::thread_rng())
            }
        };

        let n_layers = self.architecture.len() - 1;
        let n_train = x.nrows();
        let lr = F::from(config.learning_rate).ok_or_else(|| {
            StatsError::InvalidArgument(
                "learning_rate must be representable in the target float type".to_string(),
            )
        })?;

        let mut weight_ensembles: Vec<Vec<Array2<F>>> = Vec::with_capacity(config.n_ensemble);
        let mut bias_ensembles: Vec<Vec<Array1<F>>> = Vec::with_capacity(config.n_ensemble);

        for _member in 0..config.n_ensemble {
            let (x_train, y_train) = if config.bootstrap && n_train > 1 {
                let idx: Vec<usize> = (0..n_train).map(|_| rng.random_range(0..n_train)).collect();
                let xb = Array2::from_shape_fn((n_train, x.ncols()), |(i, j)| x[[idx[i], j]]);
                let yb = Array2::from_shape_fn((n_train, y.ncols()), |(i, j)| y[[idx[i], j]]);
                (xb, yb)
            } else {
                (x.to_owned(), y.to_owned())
            };

            let mut weights: Vec<Array2<F>> = Vec::with_capacity(n_layers);
            let mut biases: Vec<Array1<F>> = Vec::with_capacity(n_layers);
            for l in 0..n_layers {
                let fan_in = self.architecture[l];
                let fan_out = self.architecture[l + 1];
                let w_prec = prior_precision(&self.weight_priors[l]);
                let b_prec = prior_precision(&self.bias_priors[l]);
                let w = Array2::from_shape_fn((fan_in, fan_out), |_| {
                    sample_normal(F::zero(), w_prec, &mut rng)
                });
                let b =
                    Array1::from_shape_fn(fan_out, |_| sample_normal(F::zero(), b_prec, &mut rng));
                weights.push(w);
                biases.push(b);
            }

            for _epoch in 0..config.epochs {
                let (grads_w, grads_b) =
                    self.backward(&x_train.view(), &y_train.view(), &weights, &biases)?;
                for l in 0..n_layers {
                    let dw = grads_w[l].mapv(|g| g * lr);
                    weights[l] = &weights[l] - &dw;
                    let db = grads_b[l].mapv(|g| g * lr);
                    biases[l] = &biases[l] - &db;
                }
            }

            weight_ensembles.push(weights);
            bias_ensembles.push(biases);
        }

        self.weight_samples = Some(weight_ensembles);
        self.bias_samples = Some(bias_ensembles);
        Ok(())
    }

    /// Make predictions with uncertainty quantification.
    ///
    /// If [`Self::fit`] has already trained a posterior ensemble, the ensemble
    /// is a small *finite* population (typically a handful to a few dozen
    /// members) whose exact predictive mean/variance is directly computable,
    /// so this forward-propagates through **every** trained member exactly
    /// once (`n_samples_` is ignored in this case: randomly resampling a
    /// finite, fully-known population with replacement would only add
    /// spurious Monte Carlo noise around a quantity that has no randomness
    /// left to estimate). Otherwise -- no ensemble has been trained yet --
    /// this falls back to real prior-predictive Monte Carlo: `n_samples_`
    /// independent draws of the *whole network* from its priors, each
    /// forward-propagated through the actual input `x`, where genuine random
    /// sampling is unavoidable since the prior is a continuous distribution.
    /// Either way, the reported mean/variance are genuine empirical
    /// statistics of real forward passes -- never a constant placeholder.
    pub fn predict_with_uncertainty(
        &self,
        x: &ArrayView2<F>,
        n_samples_: usize,
    ) -> StatsResult<(Array2<F>, Array2<F>)> {
        checkarray_finite(x, "x")?;
        if n_samples_ == 0 {
            return Err(StatsError::InvalidArgument(
                "n_samples_ must be at least 1".to_string(),
            ));
        }
        if x.ncols() != self.architecture[0] {
            return Err(StatsError::DimensionMismatch(format!(
                "x has {} columns, expected {} to match the network's input layer",
                x.ncols(),
                self.architecture[0]
            )));
        }

        let n_test = x.nrows();
        let output_dim = *self.architecture.last().ok_or_else(|| {
            StatsError::InvalidArgument("architecture must be non-empty".to_string())
        })?;
        let n_layers = self.architecture.len() - 1;
        let mut rng = scirs2_core::random::thread_rng();

        let mut draws: Vec<Array2<F>> = Vec::with_capacity(n_samples_);

        match (&self.weight_samples, &self.bias_samples) {
            (Some(w_ens), Some(b_ens)) if !w_ens.is_empty() && !b_ens.is_empty() => {
                let n_members = w_ens.len().min(b_ens.len());
                for idx in 0..n_members {
                    draws.push(self.forward(x, &w_ens[idx], &b_ens[idx])?);
                }
            }
            _ => {
                for _ in 0..n_samples_ {
                    let mut weights = Vec::with_capacity(n_layers);
                    let mut biases = Vec::with_capacity(n_layers);
                    for l in 0..n_layers {
                        let fan_in = self.architecture[l];
                        let fan_out = self.architecture[l + 1];
                        let w_prec = prior_precision(&self.weight_priors[l]);
                        let b_prec = prior_precision(&self.bias_priors[l]);
                        let w = Array2::from_shape_fn((fan_in, fan_out), |_| {
                            sample_normal(F::zero(), w_prec, &mut rng)
                        });
                        let b = Array1::from_shape_fn(fan_out, |_| {
                            sample_normal(F::zero(), b_prec, &mut rng)
                        });
                        weights.push(w);
                        biases.push(b);
                    }
                    draws.push(self.forward(x, &weights, &biases)?);
                }
            }
        }

        let mut predictions = Array2::<F>::zeros((n_test, output_dim));
        let mut prediction_vars = Array2::<F>::zeros((n_test, output_dim));
        let s_f = F::from(draws.len()).expect("draw count fits in any Float");

        for i in 0..n_test {
            for j in 0..output_dim {
                let m = draws.iter().fold(F::zero(), |acc, d| acc + d[[i, j]]) / s_f;
                let v = draws
                    .iter()
                    .fold(F::zero(), |acc, d| acc + (d[[i, j]] - m) * (d[[i, j]] - m))
                    / s_f.max(F::one());
                predictions[[i, j]] = m;
                prediction_vars[[i, j]] = v;
            }
        }

        Ok((predictions, prediction_vars))
    }
}