shap_rs/explainers/
tree.rs

1use crate::{
2    explainers::ExactExplainer,
3    tree::{tree_shap, TreeEnsemble},
4    AttributionSemantics, Background, Explainer, Explanation, Predict, Result, ShapError,
5};
6use ndarray::{Array2, Array3, ArrayView2, Axis, Slice};
7
8/// Exact polynomial-time TreeSHAP for native [`TreeEnsemble`] models.
9pub struct TreeExplainer<'a> {
10    model: &'a TreeEnsemble,
11}
12impl<'a> TreeExplainer<'a> {
13    pub fn new(model: &'a TreeEnsemble) -> Self {
14        Self { model }
15    }
16    /// Explains raw outputs with per-sample base margins replacing the model's
17    /// fixed base offset. Tree contributions are unchanged; only base values shift.
18    pub fn explain_with_base_margin(
19        &self,
20        x: ArrayView2<'_, f64>,
21        base_margin: ArrayView2<'_, f64>,
22    ) -> Result<Explanation> {
23        if base_margin.dim() != (x.nrows(), self.model.n_outputs()) {
24            return Err(ShapError::DimensionMismatch {
25                expected: format!("({}, {}) base margins", x.nrows(), self.model.n_outputs()),
26                found: format!("{:?}", base_margin.dim()),
27            });
28        }
29        if base_margin.iter().any(|value| !value.is_finite()) {
30            return Err(ShapError::InvalidConfiguration(
31                "base margins must be finite".into(),
32            ));
33        }
34        let explanation = self.explain(x)?;
35        let mut bases = explanation.base_values().to_owned();
36        for row in 0..bases.nrows() {
37            for output in 0..bases.ncols() {
38                bases[[row, output]] +=
39                    base_margin[[row, output]] - self.model.base_offset()[output];
40            }
41        }
42        Explanation::new(
43            explanation.values().to_owned(),
44            bases,
45            explanation.data().to_owned(),
46        )
47        .map(|explanation| explanation.with_semantics(AttributionSemantics::TreePathDependent))
48    }
49}
50
51/// Exact interventional TreeSHAP using an explicit background distribution.
52/// Unlike [`TreeExplainer`], absent features are replaced from background rows
53/// rather than integrated using training-path covers.
54pub struct InterventionalTreeExplainer<'a> {
55    model: &'a TreeEnsemble,
56    background: Background,
57    max_features: usize,
58}
59
60impl<'a> InterventionalTreeExplainer<'a> {
61    pub fn new(model: &'a TreeEnsemble, background: Background) -> Self {
62        Self {
63            model,
64            background,
65            max_features: 20,
66        }
67    }
68
69    pub fn with_max_features(mut self, max_features: usize) -> Self {
70        self.max_features = max_features;
71        self
72    }
73
74    /// Explains probabilities exactly under the supplied background. A
75    /// one-output ensemble uses sigmoid; multiple outputs use softmax.
76    pub fn explain_probability(&self, x: ArrayView2<'_, f64>) -> Result<Explanation> {
77        ExactExplainer::new(
78            TransformedTreeModel::probability(self.model),
79            self.background.clone(),
80        )
81        .with_max_features(self.max_features)
82        .explain(x)
83        .map(|explanation| explanation.with_semantics(AttributionSemantics::Interventional))
84    }
85
86    /// Explains binary logistic loss for each sample's target label.
87    pub fn explain_binary_log_loss(
88        &self,
89        x: ArrayView2<'_, f64>,
90        targets: &[bool],
91    ) -> Result<Explanation> {
92        if self.model.n_outputs() != 1 {
93            return Err(ShapError::Unsupported(
94                "binary log-loss explanations require one raw-margin output".into(),
95            ));
96        }
97        if targets.len() != x.nrows() {
98            return Err(ShapError::DimensionMismatch {
99                expected: format!("{} binary targets", x.nrows()),
100                found: format!("{} targets", targets.len()),
101            });
102        }
103        let mut explanations = Vec::with_capacity(x.nrows());
104        for (sample, &target) in targets.iter().enumerate() {
105            explanations.push(
106                ExactExplainer::new(
107                    TransformedTreeModel::binary_log_loss(self.model, target),
108                    self.background.clone(),
109                )
110                .with_max_features(self.max_features)
111                .explain(x.slice_axis(Axis(0), Slice::from(sample..sample + 1)))?,
112            );
113        }
114        Explanation::concatenate(&explanations)
115            .map(|explanation| explanation.with_semantics(AttributionSemantics::Interventional))
116    }
117}
118
119#[derive(Clone, Copy)]
120enum TreeTransform {
121    Probability,
122    BinaryLogLoss(bool),
123}
124
125struct TransformedTreeModel<'a> {
126    model: &'a TreeEnsemble,
127    transform: TreeTransform,
128}
129
130impl<'a> TransformedTreeModel<'a> {
131    fn probability(model: &'a TreeEnsemble) -> Self {
132        Self {
133            model,
134            transform: TreeTransform::Probability,
135        }
136    }
137    fn binary_log_loss(model: &'a TreeEnsemble, target: bool) -> Self {
138        Self {
139            model,
140            transform: TreeTransform::BinaryLogLoss(target),
141        }
142    }
143}
144
145impl Predict for TransformedTreeModel<'_> {
146    fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array2<f64>> {
147        let raw = self.model.predict(x)?;
148        match self.transform {
149            TreeTransform::Probability if raw.ncols() == 1 => {
150                Ok(raw.mapv(|margin| 1.0 / (1.0 + (-margin).exp())))
151            }
152            TreeTransform::Probability => {
153                let mut probability = raw;
154                for mut row in probability.rows_mut() {
155                    let maximum = row.iter().copied().fold(f64::NEG_INFINITY, f64::max);
156                    row.mapv_inplace(|margin| (margin - maximum).exp());
157                    let total = row.sum();
158                    row.mapv_inplace(|value| value / total);
159                }
160                Ok(probability)
161            }
162            TreeTransform::BinaryLogLoss(target) => Ok(raw.mapv(|margin| {
163                margin.max(0.0) + (-margin.abs()).exp().ln_1p() - if target { margin } else { 0.0 }
164            })),
165        }
166    }
167    fn n_features(&self) -> Option<usize> {
168        Some(self.model.n_features())
169    }
170    fn n_outputs(&self) -> Option<usize> {
171        Some(self.model.n_outputs())
172    }
173}
174
175impl Explainer for InterventionalTreeExplainer<'_> {
176    fn explain(&self, x: ArrayView2<'_, f64>) -> Result<Explanation> {
177        ExactExplainer::new(self.model, self.background.clone())
178            .with_max_features(self.max_features)
179            .explain(x)
180            .map(|explanation| explanation.with_semantics(AttributionSemantics::Interventional))
181    }
182}
183impl Explainer for TreeExplainer<'_> {
184    fn explain(&self, x: ArrayView2<'_, f64>) -> Result<Explanation> {
185        let m = self.model.n_features();
186        let o = self.model.n_outputs();
187        crate::error::checked_f64_shape(&[x.nrows(), m, o], "tree explanation")?;
188        if x.nrows() == 0 {
189            return Err(ShapError::EmptyData);
190        }
191        if x.ncols() != m {
192            return Err(ShapError::DimensionMismatch {
193                expected: format!("{m} features"),
194                found: format!("{}", x.ncols()),
195            });
196        }
197        let base = self.model.expected_value();
198        let bases = Array2::from_shape_fn((x.nrows(), o), |(_, k)| base[k]);
199        let mut values = Array3::zeros((x.nrows(), m, o));
200        for i in 0..x.nrows() {
201            for (tree, weight) in self.model.trees() {
202                let phi = tree_shap(tree, x.row(i));
203                for j in 0..m {
204                    for k in 0..o {
205                        values[[i, j, k]] += weight * phi[j][k]
206                    }
207                }
208            }
209        }
210        Explanation::new(values, bases, x.to_owned())
211            .map(|explanation| explanation.with_semantics(AttributionSemantics::TreePathDependent))
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::{metrics::check_additivity, MissingBranch, Node, Predict, Tree};
219    use ndarray::array;
220
221    fn model() -> TreeEnsemble {
222        // if x0 <= 0: 1; else if x1 <= 0: 3; else: 7
223        let tree = Tree::new(
224            vec![
225                Node::Split {
226                    feature: 0,
227                    threshold: 0.0,
228                    left: 1,
229                    right: 2,
230                    missing: MissingBranch::Left,
231                    cover: 10.0,
232                },
233                Node::Leaf {
234                    values: vec![1.0],
235                    cover: 4.0,
236                },
237                Node::Split {
238                    feature: 1,
239                    threshold: 0.0,
240                    left: 3,
241                    right: 4,
242                    missing: MissingBranch::Left,
243                    cover: 6.0,
244                },
245                Node::Leaf {
246                    values: vec![3.0],
247                    cover: 2.0,
248                },
249                Node::Leaf {
250                    values: vec![7.0],
251                    cover: 4.0,
252                },
253            ],
254            0,
255            2,
256        )
257        .unwrap();
258        TreeEnsemble::new(vec![(tree, 1.0)], vec![0.5]).unwrap()
259    }
260
261    #[test]
262    fn tree_shap_is_additive() {
263        let model = model();
264        let x = array![[1.0, 1.0], [-1.0, 9.0]];
265        let explanation = TreeExplainer::new(&model).explain(x.view()).unwrap();
266        let prediction = model.predict(x.view()).unwrap();
267        check_additivity(&explanation, prediction.view(), 1e-10).unwrap();
268        assert!((explanation.base_values()[[0, 0]] - 4.3).abs() < 1e-12);
269    }
270
271    #[test]
272    fn missing_values_follow_configured_branch() {
273        let model = model();
274        let prediction = model.predict(array![[f64::NAN, 2.0]].view()).unwrap();
275        assert_eq!(prediction[[0, 0]], 1.5);
276    }
277
278    #[test]
279    fn base_margins_replace_the_fixed_offset() {
280        let model = model();
281        let x = array![[1.0, 1.0], [-1.0, 9.0]];
282        let margins = array![[2.0], [-3.0]];
283        let prediction = model
284            .predict_with_base_margin(x.view(), margins.view())
285            .unwrap();
286        let explanation = TreeExplainer::new(&model)
287            .explain_with_base_margin(x.view(), margins.view())
288            .unwrap();
289        check_additivity(&explanation, prediction.view(), 1e-10).unwrap();
290        assert!(model
291            .predict_with_base_margin(x.view(), array![[1.0, 2.0]].view())
292            .is_err());
293    }
294
295    #[test]
296    fn interventional_tree_values_use_the_supplied_background() {
297        let model = model();
298        let background = Background::new(array![[-1.0, -1.0], [1.0, 1.0]]).unwrap();
299        let x = array![[1.0, -1.0]];
300        let interventional = InterventionalTreeExplainer::new(&model, background.clone())
301            .explain(x.view())
302            .unwrap();
303        let exact = ExactExplainer::new(&model, background)
304            .explain(x.view())
305            .unwrap();
306        assert_eq!(interventional.values(), exact.values());
307        assert_eq!(interventional.base_values(), exact.base_values());
308        check_additivity(
309            &interventional,
310            model.predict(x.view()).unwrap().view(),
311            1e-10,
312        )
313        .unwrap();
314    }
315
316    #[test]
317    fn interventional_probability_and_log_loss_are_additive() {
318        let model = model();
319        let background = Background::new(array![[-1.0, -1.0], [1.0, 1.0]]).unwrap();
320        let x = array![[1.0, -1.0], [-1.0, 2.0]];
321        let explainer = InterventionalTreeExplainer::new(&model, background);
322        let probability = explainer.explain_probability(x.view()).unwrap();
323        let expected_probability = TransformedTreeModel::probability(&model)
324            .predict(x.view())
325            .unwrap();
326        check_additivity(&probability, expected_probability.view(), 1e-10).unwrap();
327
328        let loss = explainer
329            .explain_binary_log_loss(x.view(), &[true, false])
330            .unwrap();
331        for sample in 0..x.nrows() {
332            let expected = TransformedTreeModel::binary_log_loss(&model, sample == 0)
333                .predict(x.slice_axis(Axis(0), Slice::from(sample..sample + 1)))
334                .unwrap();
335            assert!((loss.reconstructed()[[sample, 0]] - expected[[0, 0]]).abs() < 1e-10);
336        }
337    }
338}
shap_rs/explainers/tree.rs

shap_rs/explainers/
tree.rs
