use std::collections::BTreeMap;
use crate::approx::approximate_asv;
use crate::error::CausasvError;
use crate::graph::{Dag, NodeId};
use crate::sampler::SamplingConfig;
use crate::topo::enumerate_topos;
use crate::tree::tree_exact_asv;
pub struct AsvResult {
pub values: BTreeMap<NodeId, f64>,
pub n_samples: usize,
pub seed: Option<u64>,
pub is_exact: bool,
pub effective_sample_size: Option<f64>,
}
pub struct AsvExplainer {
dag: Dag,
}
impl AsvExplainer {
pub fn new(dag: Dag) -> Self {
Self { dag }
}
pub fn exact<F>(&self, value_fn: F) -> Result<AsvResult, CausasvError>
where
F: Fn(&[NodeId]) -> Result<f64, CausasvError>,
{
self.dag.validate()?;
let orderings = enumerate_topos(&self.dag)?;
let n_orderings = orderings.len();
let n = self.dag.node_count();
let mut phi = vec![0.0f64; n];
let mut coalition = Vec::with_capacity(n);
for ordering in &orderings {
for (k, &node) in ordering.iter().enumerate() {
coalition.clear();
coalition.extend_from_slice(&ordering[..k]);
coalition.sort_unstable();
let v_without = value_fn(&coalition)?;
coalition.push(node);
coalition.sort_unstable();
let v_with = value_fn(&coalition)?;
phi[node.0 as usize] += v_with - v_without;
}
}
let scale = 1.0 / n_orderings as f64;
let values = (0..n).map(|i| (NodeId(i as u32), phi[i] * scale)).collect();
Ok(AsvResult {
values,
n_samples: n_orderings,
seed: None,
is_exact: true,
effective_sample_size: None,
})
}
pub fn approximate<F>(
&self,
value_fn: F,
config: SamplingConfig,
) -> Result<AsvResult, CausasvError>
where
F: Fn(&[NodeId]) -> Result<f64, CausasvError>,
{
self.dag.validate()?;
approximate_asv(&self.dag, value_fn, config)
}
pub fn exact_tree<F>(&self, value_fn: F) -> Result<AsvResult, CausasvError>
where
F: Fn(&[NodeId]) -> Result<f64, CausasvError>,
{
self.dag.validate()?;
tree_exact_asv(&self.dag, value_fn)
}
}