causasv 0.2.0

Fast causal Asymmetric Shapley Values for Rust and Python
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use std::collections::HashMap;

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;

use crate::asv::AsvExplainer;
use crate::error::CausasvError;
use crate::graph::{Dag as RustDag, NodeId};
use crate::sampler::SamplingConfig;

#[pyclass(name = "CausalDAG")]
pub struct PyCausalDAG {
    pub(crate) inner: RustDag,
}

#[pymethods]
impl PyCausalDAG {
    #[new]
    fn new() -> Self {
        Self {
            inner: RustDag::new(),
        }
    }

    /// Add a directed edge. Nodes are created automatically if they do not exist.
    fn add_edge(&mut self, from_name: &str, to_name: &str) -> PyResult<()> {
        let from = self.inner.add_node(from_name);
        let to = self.inner.add_node(to_name);
        self.inner.add_edge(from, to).map_err(py_err)
    }

    /// Validate the graph (check for cycles, empty graph, etc.).
    fn validate(&self) -> PyResult<()> {
        self.inner.validate().map_err(py_err)
    }
}

#[pyclass(name = "ASVExplainer")]
pub struct PyASVExplainer {
    inner: AsvExplainer,
    /// NodeId index → node name, used to bridge coalitions to Python strings.
    names: Vec<String>,
}

#[pymethods]
impl PyASVExplainer {
    /// Create an explainer from a CausalDAG. Validates the graph immediately.
    #[new]
    fn new(dag: &PyCausalDAG) -> PyResult<Self> {
        dag.inner.validate().map_err(py_err)?;
        let names: Vec<String> = dag
            .inner
            .all_nodes()
            .map(|id| dag.inner.node_name(id).unwrap().to_string())
            .collect();
        Ok(Self {
            inner: AsvExplainer::new(dag.inner.clone()),
            names,
        })
    }

    /// Compute ASV values.
    ///
    /// Args:
    ///   value_fn: callable (list[str]) -> float.
    ///             Receives a sorted list of node names in the coalition.
    ///   method:   "approx" (default), "exact", or "exact_tree".
    ///   n_samples: number of samples for approximate method (default 10_000).
    ///   seed:     RNG seed for reproducibility (default None = random).
    ///
    /// Returns:
    ///   dict[str, float] mapping node name to its ASV value.
    #[pyo3(signature = (value_fn, method = "approx", n_samples = 10_000, seed = None))]
    fn explain(
        &self,
        py: Python<'_>,
        value_fn: PyObject,
        method: &str,
        n_samples: usize,
        seed: Option<u64>,
    ) -> PyResult<HashMap<String, f64>> {
        let names = &self.names;
        let rust_fn = |coalition: &[NodeId]| -> Result<f64, CausasvError> {
            let name_list: Vec<&str> = coalition
                .iter()
                .map(|id| names[id.0 as usize].as_str())
                .collect();
            value_fn
                .call1(py, (name_list,))
                .and_then(|r| r.extract::<f64>(py))
                .map_err(|e| CausasvError::ValueFunctionError(e.to_string()))
        };

        let result = match method {
            "approx" => {
                let mut cfg = SamplingConfig::new(n_samples);
                if let Some(s) = seed {
                    cfg = cfg.with_seed(s);
                }
                self.inner.approximate(rust_fn, cfg)
            }
            "exact" => self.inner.exact(rust_fn),
            "exact_tree" => self.inner.exact_tree(rust_fn),
            _ => {
                return Err(PyValueError::new_err(format!(
                    "unknown method '{method}': use 'approx', 'exact', or 'exact_tree'"
                )))
            }
        }
        .map_err(py_err)?;

        Ok(result
            .values
            .iter()
            .map(|(id, &v)| (names[id.0 as usize].clone(), v))
            .collect())
    }
}

#[pymodule]
fn causasv(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyCausalDAG>()?;
    m.add_class::<PyASVExplainer>()?;
    Ok(())
}

fn py_err(e: CausasvError) -> PyErr {
    PyValueError::new_err(e.to_string())
}