shap_rs/explainers/

gradient.rs

use crate::{
    Background, DifferentiablePredict, Explainer, Explanation, Result, ShapError,
    UncertainExplanation,
};
use ndarray::{Array2, Array3, ArrayView2, Axis, Slice};
use rand::{rngs::StdRng, Rng, SeedableRng};
/// Expected Gradients (Gradient SHAP) with background interpolation and
/// optional Gaussian local smoothing.
pub struct GradientExplainer<M> {
    model: M,
    background: Background,
    nsamples: usize,
    seed: u64,
    local_smoothing: f64,
    batch_size: usize,
}
impl<M> GradientExplainer<M> {
    pub fn new(model: M, background: Background) -> Self {
        Self {
            model,
            background,
            nsamples: 256,
            seed: 0,
            local_smoothing: 0.0,
            batch_size: 256,
        }
    }
    pub fn with_nsamples(mut self, n: usize) -> Self {
        self.nsamples = n;
        self
    }
    pub fn with_seed(mut self, s: u64) -> Self {
        self.seed = s;
        self
    }
    pub fn with_local_smoothing(mut self, s: f64) -> Self {
        self.local_smoothing = s;
        self
    }
    /// Limits gradient rows submitted to the autodiff backend in one call.
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }
}
impl<M: DifferentiablePredict> GradientExplainer<M> {
    /// Repeats Expected Gradients with independent deterministic seeds and
    /// returns the standard error of the mean attribution.
    pub fn explain_with_uncertainty(
        &self,
        x: ArrayView2<'_, f64>,
        repeats: usize,
    ) -> Result<UncertainExplanation> {
        if repeats < 2 {
            return Err(ShapError::InvalidConfiguration(
                "uncertainty estimation requires at least two repeats".into(),
            ));
        }
        let mut runs = Vec::with_capacity(repeats);
        for repeat in 0..repeats {
            runs.push(
                GradientExplainer {
                    model: &self.model,
                    background: self.background.clone(),
                    nsamples: self.nsamples,
                    seed: self.seed.wrapping_add(repeat as u64),
                    local_smoothing: self.local_smoothing,
                    batch_size: self.batch_size,
                }
                .explain(x)?,
            );
        }
        let shape = runs[0].values().dim();
        let mut mean = Array3::<f64>::zeros(shape);
        for run in &runs {
            ndarray::Zip::from(&mut mean)
                .and(run.values())
                .for_each(|average, &value| *average += value);
        }
        mean.mapv_inplace(|value| value / repeats as f64);
        let mut variance = Array3::<f64>::zeros(shape);
        for run in &runs {
            ndarray::Zip::from(&mut variance)
                .and(run.values())
                .and(&mean)
                .for_each(|sum, &value, &average| *sum += (value - average).powi(2));
        }
        let standard_errors =
            variance.mapv(|value| (value / ((repeats - 1) * repeats) as f64).sqrt());
        let explanation = Explanation::new(mean, runs[0].base_values().to_owned(), x.to_owned())?;
        UncertainExplanation::new(explanation, standard_errors, repeats)
    }
}
impl<M: DifferentiablePredict> Explainer for GradientExplainer<M> {
    fn explain(&self, x: ArrayView2<'_, f64>) -> Result<Explanation> {
        let m = self.background.n_features();
        if x.nrows() == 0 {
            return Err(ShapError::EmptyData);
        }
        if x.ncols() != m {
            return Err(ShapError::DimensionMismatch {
                expected: format!("{m} features"),
                found: format!("{}", x.ncols()),
            });
        }
        if self.nsamples == 0
            || self.batch_size == 0
            || !self.local_smoothing.is_finite()
            || self.local_smoothing < 0.
        {
            return Err(ShapError::InvalidConfiguration(
                "nsamples and batch size must be positive and local smoothing non-negative".into(),
            ));
        }
        let prediction = self.model.predict(self.background.data())?;
        if prediction.nrows() != self.background.n_samples() || prediction.ncols() == 0 {
            return Err(ShapError::DimensionMismatch {
                expected: format!("{} background predictions", self.background.n_samples()),
                found: format!("{:?}", prediction.dim()),
            });
        }
        let base = prediction.mean_axis(Axis(0)).unwrap();
        let o = base.len();
        crate::error::checked_f64_shape(&[x.nrows(), m, o], "gradient explanation")?;
        crate::error::checked_f64_shape(&[self.nsamples, m], "gradient sampling batch")?;
        let bases = Array2::from_shape_fn((x.nrows(), o), |(_, k)| base[k]);
        let std = feature_std(&self.background);
        let mut values = Array3::zeros((x.nrows(), m, o));
        for n in 0..x.nrows() {
            let mut rng = StdRng::seed_from_u64(crate::coalition::sample_seed(self.seed, x.row(n)));
            let mut points = Array2::zeros((self.nsamples, m));
            let mut deltas = Array2::zeros((self.nsamples, m));
            for s in 0..self.nsamples {
                let b = rng.gen_range(0..self.background.n_samples());
                let alpha = rng.gen::<f64>();
                for j in 0..m {
                    let noise = if self.local_smoothing > 0. {
                        gaussian(&mut rng) * self.local_smoothing * std[j]
                    } else {
                        0.
                    };
                    let delta = x[[n, j]] + noise - self.background.data()[[b, j]];
                    deltas[[s, j]] = delta;
                    points[[s, j]] = self.background.data()[[b, j]] + alpha * delta
                }
            }
            for start in (0..self.nsamples).step_by(self.batch_size) {
                let end = start.saturating_add(self.batch_size).min(self.nsamples);
                let gradients = self
                    .model
                    .gradients(points.slice_axis(Axis(0), Slice::from(start..end)))?;
                if gradients.dim() != (end - start, m, o) {
                    return Err(ShapError::DimensionMismatch {
                        expected: format!("({}, {m}, {o}) gradients", end - start),
                        found: format!("{:?}", gradients.dim()),
                    });
                }
                if gradients.iter().any(|v| !v.is_finite()) {
                    return Err(ShapError::ModelError(
                        "gradient contains a non-finite value".into(),
                    ));
                }
                for j in 0..m {
                    for k in 0..o {
                        values[[n, j, k]] += (start..end)
                            .map(|sample| gradients[[sample - start, j, k]] * deltas[[sample, j]])
                            .sum::<f64>()
                            / self.nsamples as f64
                    }
                }
            }
        }
        Explanation::new(values, bases, x.to_owned())
    }
}
fn gaussian<R: Rng + ?Sized>(rng: &mut R) -> f64 {
    let u1 = rng.gen::<f64>().max(f64::MIN_POSITIVE);
    let u2 = rng.gen::<f64>();
    (-2. * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
}
fn feature_std(bg: &Background) -> Vec<f64> {
    let mean = bg.data().mean_axis(Axis(0)).unwrap();
    (0..bg.n_features())
        .map(|j| {
            (bg.data()
                .column(j)
                .iter()
                .map(|x| (x - mean[j]).powi(2))
                .sum::<f64>()
                / bg.n_samples() as f64)
                .sqrt()
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Predict;
    use ndarray::array;
    struct Linear;
    impl Predict for Linear {
        fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array2<f64>> {
            Ok(x.map_axis(Axis(1), |r| 2. * r[0] - r[1])
                .insert_axis(Axis(1)))
        }
    }
    impl DifferentiablePredict for Linear {
        fn gradients(&self, x: ArrayView2<'_, f64>) -> Result<Array3<f64>> {
            Ok(Array3::from_shape_fn((x.nrows(), 2, 1), |(_, j, _)| {
                if j == 0 {
                    2.
                } else {
                    -1.
                }
            }))
        }
    }
    #[test]
    fn expected_gradients_is_exact_for_linear_models() {
        let e = GradientExplainer::new(Linear, Background::new(array![[0., 0.]]).unwrap())
            .with_nsamples(32)
            .with_batch_size(3)
            .explain(array![[3., 4.]].view())
            .unwrap();
        assert!((e.values()[[0, 0, 0]] - 6.).abs() < 1e-12);
        assert!((e.values()[[0, 1, 0]] + 4.).abs() < 1e-12);
        assert!((e.reconstructed()[[0, 0]] - 2.).abs() < 1e-12);
    }

    #[test]
    fn reports_uncertainty_for_stochastic_expected_gradients() {
        let e =
            GradientExplainer::new(Linear, Background::new(array![[0., 0.], [2., 4.]]).unwrap())
                .with_nsamples(16)
                .explain_with_uncertainty(array![[3., 4.]].view(), 4)
                .unwrap();
        assert_eq!(e.repeats(), 4);
        assert_eq!(e.standard_errors().dim(), (1, 2, 1));
        assert!(e.standard_errors().iter().all(|value| value.is_finite()));
    }
}