Skip to main content

shap_rs/explainers/
causal.rs

1use crate::{
2    evaluation::CoalitionEvaluator, AttributionSemantics, Background, ConditionalTabularMasker,
3    EvaluationConfig, Explainer, Explanation, IndependentMasker, Link, Masker, Predict, Result,
4    ShapError,
5};
6use ndarray::{Array2, Array3, ArrayView1, ArrayView2};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum CausalMaskingMode {
10    Interventional,
11    Observational,
12    Conditional,
13}
14
15#[derive(Debug, Clone)]
16pub struct CausalTabularMasker {
17    mode: CausalMaskingMode,
18    interventional: Option<IndependentMasker>,
19    conditional: Option<ConditionalTabularMasker>,
20}
21
22impl CausalTabularMasker {
23    pub fn interventional(background: Background) -> Self {
24        Self {
25            mode: CausalMaskingMode::Interventional,
26            interventional: Some(IndependentMasker::new(background)),
27            conditional: None,
28        }
29    }
30    pub fn observational(
31        background: Background,
32        categorical_features: &[usize],
33        neighbors: usize,
34    ) -> Result<Self> {
35        Self::nearest(
36            CausalMaskingMode::Observational,
37            background,
38            categorical_features,
39            neighbors,
40        )
41    }
42    pub fn conditional(
43        background: Background,
44        categorical_features: &[usize],
45        neighbors: usize,
46    ) -> Result<Self> {
47        Self::nearest(
48            CausalMaskingMode::Conditional,
49            background,
50            categorical_features,
51            neighbors,
52        )
53    }
54    fn nearest(
55        mode: CausalMaskingMode,
56        background: Background,
57        categorical_features: &[usize],
58        neighbors: usize,
59    ) -> Result<Self> {
60        Ok(Self {
61            mode,
62            interventional: None,
63            conditional: Some(ConditionalTabularMasker::new(
64                background,
65                categorical_features,
66                neighbors,
67            )?),
68        })
69    }
70    pub fn mode(&self) -> CausalMaskingMode {
71        self.mode
72    }
73}
74
75impl Masker for CausalTabularMasker {
76    fn n_features(&self) -> usize {
77        self.interventional
78            .as_ref()
79            .map(Masker::n_features)
80            .or_else(|| self.conditional.as_ref().map(Masker::n_features))
81            .unwrap_or(0)
82    }
83    fn mask(&self, sample: ArrayView1<'_, f64>, present: &[bool]) -> Result<Array2<f64>> {
84        if let Some(masker) = &self.interventional {
85            masker.mask(sample, present)
86        } else if let Some(masker) = &self.conditional {
87            masker.mask(sample, present)
88        } else {
89            Err(ShapError::InvalidConfiguration(
90                "causal tabular masker has no sampling strategy".into(),
91            ))
92        }
93    }
94}
95
96/// Directed acyclic dependency graph for asymmetric causal Shapley values.
97#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
98#[serde(try_from = "CausalGraphPayload")]
99pub struct CausalGraph {
100    parents: Vec<Vec<usize>>,
101}
102#[derive(serde::Deserialize)]
103struct CausalGraphPayload {
104    parents: Vec<Vec<usize>>,
105}
106impl TryFrom<CausalGraphPayload> for CausalGraph {
107    type Error = ShapError;
108    fn try_from(payload: CausalGraphPayload) -> Result<Self> {
109        Self::new(payload.parents)
110    }
111}
112impl CausalGraph {
113    pub fn new(parents: Vec<Vec<usize>>) -> Result<Self> {
114        let n = parents.len();
115        if n == 0 {
116            return Err(ShapError::InvalidConfiguration(
117                "causal graph cannot be empty".into(),
118            ));
119        }
120        for (i, p) in parents.iter().enumerate() {
121            let mut q = p.clone();
122            q.sort();
123            q.dedup();
124            if q.len() != p.len() || p.iter().any(|&j| j >= n || j == i) {
125                return Err(ShapError::InvalidConfiguration(
126                    "causal graph has invalid or duplicate parents".into(),
127                ));
128            }
129        }
130        let graph = Self { parents };
131        if !graph.is_acyclic() {
132            return Err(ShapError::InvalidConfiguration(
133                "causal graph contains a cycle".into(),
134            ));
135        }
136        Ok(graph)
137    }
138    pub fn parents(&self) -> &[Vec<usize>] {
139        &self.parents
140    }
141    pub fn n_features(&self) -> usize {
142        self.parents.len()
143    }
144    /// Revalidates parent indices, uniqueness, and acyclicity after deserialization.
145    pub fn validate(&self) -> Result<()> {
146        Self::new(self.parents.clone()).map(|_| ())
147    }
148    fn is_acyclic(&self) -> bool {
149        let mut used = vec![false; self.n_features()];
150        for _ in 0..self.n_features() {
151            if let Some(j) = (0..self.n_features())
152                .find(|&j| !used[j] && self.parents[j].iter().all(|&p| used[p]))
153            {
154                used[j] = true
155            } else {
156                return false;
157            }
158        }
159        true
160    }
161    fn topological_orders(&self, limit: usize) -> Result<Vec<Vec<usize>>> {
162        fn rec(
163            g: &CausalGraph,
164            used: &mut [bool],
165            order: &mut Vec<usize>,
166            out: &mut Vec<Vec<usize>>,
167            limit: usize,
168        ) {
169            if out.len() > limit {
170                return;
171            }
172            if order.len() == used.len() {
173                out.push(order.clone());
174                return;
175            }
176            for j in 0..used.len() {
177                if !used[j] && g.parents[j].iter().all(|&p| used[p]) {
178                    used[j] = true;
179                    order.push(j);
180                    rec(g, used, order, out, limit);
181                    order.pop();
182                    used[j] = false
183                }
184            }
185        }
186        let mut out = Vec::new();
187        rec(
188            self,
189            &mut vec![false; self.n_features()],
190            &mut Vec::new(),
191            &mut out,
192            limit,
193        );
194        if out.len() > limit {
195            return Err(ShapError::InvalidConfiguration(format!(
196                "causal graph generates more than {limit} topological orders"
197            )));
198        }
199        Ok(out)
200    }
201}
202
203/// Asymmetric causal Shapley values averaged over all topological orderings.
204pub struct CausalExplainer<M, K> {
205    model: M,
206    masker: K,
207    graph: CausalGraph,
208    max_orders: usize,
209    evaluation: EvaluationConfig,
210    link: Link,
211}
212impl<M, K> CausalExplainer<M, K> {
213    pub fn new(model: M, masker: K, graph: CausalGraph) -> Self {
214        Self {
215            model,
216            masker,
217            graph,
218            max_orders: 65536,
219            evaluation: EvaluationConfig {
220                coalition_batch_size: 64,
221                cache_capacity: 1 << 20,
222                max_model_rows: None,
223            },
224            link: Link::Identity,
225        }
226    }
227    pub fn with_max_orders(mut self, n: usize) -> Self {
228        self.max_orders = n;
229        self
230    }
231    pub fn with_evaluation_config(mut self, c: EvaluationConfig) -> Self {
232        self.evaluation = c;
233        self
234    }
235    pub fn with_link(mut self, link: Link) -> Self {
236        self.link = link;
237        self
238    }
239}
240impl<M: Predict, K: Masker> Explainer for CausalExplainer<M, K> {
241    fn explain(&self, x: ArrayView2<'_, f64>) -> Result<Explanation> {
242        let m = self.masker.n_features();
243        self.graph.validate()?;
244        if x.nrows() == 0 {
245            return Err(ShapError::EmptyData);
246        }
247        if x.ncols() != m || self.graph.n_features() != m {
248            return Err(ShapError::DimensionMismatch {
249                expected: format!("{m} features in data and causal graph"),
250                found: format!("data {}, graph {}", x.ncols(), self.graph.n_features()),
251            });
252        }
253        if m >= 63 {
254            return Err(ShapError::InvalidConfiguration(
255                "causal explanations support at most 62 features".into(),
256            ));
257        }
258        let orders = self.graph.topological_orders(self.max_orders)?;
259        let step_count = orders.len().checked_mul(m).ok_or_else(|| {
260            ShapError::InvalidConfiguration("causal order step count overflowed".into())
261        })?;
262        crate::error::checked_f64_shape(&[step_count], "causal order steps")?;
263        let mut probe = CoalitionEvaluator::new(&self.model, &self.masker, self.evaluation)?;
264        let o = probe.evaluate(x.row(0), &[0])?[0].len();
265        crate::error::checked_f64_shape(&[x.nrows(), m, o], "causal explanation")?;
266        let mut values = Array3::zeros((x.nrows(), m, o));
267        let mut bases = Array2::zeros((x.nrows(), o));
268        for n in 0..x.nrows() {
269            let mut masks = vec![0u64];
270            let mut steps = Vec::with_capacity(step_count);
271            for order in &orders {
272                let mut mask = 0;
273                let mut before = 0;
274                for &j in order {
275                    mask |= 1 << j;
276                    masks.push(mask);
277                    let after = masks.len() - 1;
278                    steps.push((j, before, after));
279                    before = after
280                }
281            }
282            let mut evaluator =
283                CoalitionEvaluator::new(&self.model, &self.masker, self.evaluation)?;
284            let evaluated = evaluator
285                .evaluate(x.row(n), &masks)?
286                .into_iter()
287                .map(|row| {
288                    row.into_iter()
289                        .map(|value| self.link.forward(value))
290                        .collect::<Result<Vec<_>>>()
291                })
292                .collect::<Result<Vec<_>>>()?;
293            for k in 0..o {
294                bases[[n, k]] = evaluated[0][k]
295            }
296            for (j, before, after) in steps {
297                for k in 0..o {
298                    values[[n, j, k]] +=
299                        (evaluated[after][k] - evaluated[before][k]) / orders.len() as f64
300                }
301            }
302        }
303        Explanation::new(values, bases, x.to_owned())
304            .map(|explanation| explanation.with_semantics(AttributionSemantics::CausalAsymmetric))
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311    use crate::{FixedMasker, FnModel};
312    use ndarray::{array, Axis};
313    #[test]
314    fn causal_order_allocates_interaction_asymmetrically() {
315        let graph = CausalGraph::new(vec![vec![], vec![0]]).unwrap();
316        let model = FnModel::new(|x: ArrayView2<'_, f64>| {
317            Ok(x.map_axis(Axis(1), |r| r[0] * r[1]).insert_axis(Axis(1)))
318        });
319        let e = CausalExplainer::new(model, FixedMasker::new(array![0., 0.]).unwrap(), graph)
320            .explain(array![[1., 1.]].view())
321            .unwrap();
322        assert!(e.values()[[0, 0, 0]].abs() < 1e-12);
323        assert!((e.values()[[0, 1, 0]] - 1.).abs() < 1e-12);
324        assert_eq!(e.semantics(), AttributionSemantics::CausalAsymmetric);
325    }
326    #[test]
327    fn causal_tabular_modes_select_interventional_or_conditional_sampling() {
328        let background = Background::new(array![[0., 0.], [1., 0.1], [1., 9.]]).unwrap();
329        let interventional = CausalTabularMasker::interventional(background.clone());
330        assert_eq!(
331            interventional
332                .mask(array![1., 8.].view(), &[true, false])
333                .unwrap()
334                .nrows(),
335            3
336        );
337        let observational = CausalTabularMasker::observational(background, &[0], 1).unwrap();
338        assert_eq!(observational.mode(), CausalMaskingMode::Observational);
339        assert_eq!(
340            observational
341                .mask(array![1., 8.].view(), &[true, false])
342                .unwrap(),
343            array![[1., 0.1]]
344        );
345    }
346    #[test]
347    fn rejects_invalid_deserialized_style_graph_before_evaluation() {
348        let graph = CausalGraph {
349            parents: vec![vec![1], vec![0]],
350        };
351        let model =
352            FnModel::new(|x: ArrayView2<'_, f64>| Ok(x.sum_axis(Axis(1)).insert_axis(Axis(1))));
353        let result = CausalExplainer::new(model, FixedMasker::new(array![0., 0.]).unwrap(), graph)
354            .explain(array![[1., 1.]].view());
355        assert!(matches!(result, Err(ShapError::InvalidConfiguration(_))));
356    }
357    #[test]
358    fn causal_logit_link_explains_log_odds() {
359        let graph = CausalGraph::new(vec![vec![]]).unwrap();
360        let model =
361            FnModel::new(|x: ArrayView2<'_, f64>| Ok(x.column(0).to_owned().insert_axis(Axis(1))));
362        let explanation =
363            CausalExplainer::new(model, FixedMasker::new(array![0.5]).unwrap(), graph)
364                .with_link(Link::Logit)
365                .explain(array![[0.8]].view())
366                .unwrap();
367        assert!((explanation.reconstructed()[[0, 0]] - 4f64.ln()).abs() < 1e-12);
368    }
369}