1use crate::{Background, DeepAttribution, Explainer, Explanation, Result, ShapError};
2use ndarray::{Array2, ArrayView2, Axis, Slice};
3pub 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 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}