Skip to main content

shap_rs/explainers/
deep.rs

1use crate::{Background, DeepAttribution, Explainer, Explanation, Result, ShapError};
2use ndarray::{Array2, ArrayView2, Axis, Slice};
3/// Deep SHAP coordinator for framework adapters implementing multiplier
4/// propagation through their native computation graph.
5pub struct DeepExplainer<M> {
6    model: M,
7    background: Background,
8    check_additivity: bool,
9    tolerance: f64,
10    batch_size: usize,
11}
12impl<M> DeepExplainer<M> {
13    pub fn new(model: M, background: Background) -> Self {
14        Self {
15            model,
16            background,
17            check_additivity: true,
18            tolerance: 1e-5,
19            batch_size: 256,
20        }
21    }
22    pub fn with_additivity_check(mut self, enabled: bool) -> Self {
23        self.check_additivity = enabled;
24        self
25    }
26    pub fn with_tolerance(mut self, t: f64) -> Self {
27        self.tolerance = t;
28        self
29    }
30    /// Limits explained samples submitted to the graph adapter in one call.
31    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
32        self.batch_size = batch_size;
33        self
34    }
35}
36impl<M: DeepAttribution> Explainer for DeepExplainer<M> {
37    fn explain(&self, x: ArrayView2<'_, f64>) -> Result<Explanation> {
38        if !self.tolerance.is_finite() || self.tolerance < 0.0 || self.batch_size == 0 {
39            return Err(ShapError::InvalidConfiguration(
40                "Deep SHAP tolerance must be finite and non-negative and batch size positive"
41                    .into(),
42            ));
43        }
44        if x.nrows() == 0 {
45            return Err(ShapError::EmptyData);
46        }
47        if x.ncols() != self.background.n_features() {
48            return Err(ShapError::DimensionMismatch {
49                expected: format!("{} features", self.background.n_features()),
50                found: format!("{}", x.ncols()),
51            });
52        }
53        if self
54            .model
55            .n_features()
56            .is_some_and(|features| features != x.ncols())
57        {
58            return Err(ShapError::DimensionMismatch {
59                expected: format!("model with {} features", x.ncols()),
60                found: format!("model reports {:?}", self.model.n_features()),
61            });
62        }
63        let bg_pred = self.model.predict(self.background.data())?;
64        if bg_pred.nrows() != self.background.n_samples() || bg_pred.ncols() == 0 {
65            return Err(ShapError::DimensionMismatch {
66                expected: format!("{} background predictions", self.background.n_samples()),
67                found: format!("{:?}", bg_pred.dim()),
68            });
69        }
70        let base = bg_pred.mean_axis(Axis(0)).unwrap();
71        crate::error::checked_f64_shape(&[x.nrows(), x.ncols(), base.len()], "deep explanation")?;
72        if self
73            .model
74            .n_outputs()
75            .is_some_and(|outputs| outputs != base.len())
76        {
77            return Err(ShapError::OutputDimensionMismatch {
78                expected: self.model.n_outputs().unwrap(),
79                found: base.len(),
80            });
81        }
82        let bases = Array2::from_shape_fn((x.nrows(), base.len()), |(_, o)| base[o]);
83        let mut parts = Vec::new();
84        for start in (0..x.nrows()).step_by(self.batch_size) {
85            let end = start.saturating_add(self.batch_size).min(x.nrows());
86            parts.push(self.model.deep_contributions(
87                x.slice_axis(Axis(0), Slice::from(start..end)),
88                self.background.data(),
89            )?);
90        }
91        let views = parts.iter().map(|part| part.view()).collect::<Vec<_>>();
92        let values = ndarray::concatenate(Axis(0), &views)
93            .map_err(|error| ShapError::ModelError(error.to_string()))?;
94        let e = Explanation::new(values, bases, x.to_owned())?;
95        if self.check_additivity {
96            let mut predictions = Vec::new();
97            for start in (0..x.nrows()).step_by(self.batch_size) {
98                let end = start.saturating_add(self.batch_size).min(x.nrows());
99                predictions.push(
100                    self.model
101                        .predict(x.slice_axis(Axis(0), Slice::from(start..end)))?,
102                );
103            }
104            let views = predictions
105                .iter()
106                .map(|part| part.view())
107                .collect::<Vec<_>>();
108            let prediction = ndarray::concatenate(Axis(0), &views)
109                .map_err(|error| ShapError::ModelError(error.to_string()))?;
110            crate::metrics::check_additivity(&e, prediction.view(), self.tolerance)?
111        }
112        Ok(e)
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use crate::{DeepAttribution, Predict};
120    use ndarray::{array, Array3};
121    struct Adapter;
122    impl Predict for Adapter {
123        fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array2<f64>> {
124            Ok(x.sum_axis(Axis(1)).insert_axis(Axis(1)))
125        }
126    }
127    impl DeepAttribution for Adapter {
128        fn deep_contributions(
129            &self,
130            x: ArrayView2<'_, f64>,
131            bg: ArrayView2<'_, f64>,
132        ) -> Result<Array3<f64>> {
133            let mean = bg.mean_axis(Axis(0)).unwrap();
134            Ok(Array3::from_shape_fn(
135                (x.nrows(), x.ncols(), 1),
136                |(i, j, _)| x[[i, j]] - mean[j],
137            ))
138        }
139    }
140    #[test]
141    fn validates_adapter_contributions() {
142        let e = DeepExplainer::new(
143            Adapter,
144            Background::new(array![[0., 0.], [2., 2.]]).unwrap(),
145        )
146        .explain(array![[3., 4.]].view())
147        .unwrap();
148        assert_eq!(e.base_values()[[0, 0]], 2.);
149        assert_eq!(e.reconstructed()[[0, 0]], 7.);
150    }
151
152    #[test]
153    fn rejects_invalid_tolerance() {
154        let result = DeepExplainer::new(Adapter, Background::new(array![[0., 0.]]).unwrap())
155            .with_tolerance(f64::NAN)
156            .explain(array![[1., 1.]].view());
157        assert!(matches!(result, Err(ShapError::InvalidConfiguration(_))));
158    }
159
160    #[test]
161    fn mini_batches_contributions_and_predictions() {
162        let e = DeepExplainer::new(Adapter, Background::new(array![[0., 0.]]).unwrap())
163            .with_batch_size(1)
164            .explain(array![[1., 2.], [3., 4.], [5., 6.]].view())
165            .unwrap();
166        assert_eq!(e.values().dim(), (3, 2, 1));
167        assert_eq!(e.reconstructed()[[2, 0]], 11.0);
168    }
169}