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#[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 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
203pub 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}