causal-hub 0.0.4

A library for causal models, inference and discovery.
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
130
131

use dry::macro_for;
use ndarray::prelude::*;

use crate::{
    datasets::{CatIncTable, CatTable, CatWtdTable},
    estimators::{BE, CPDEstimator, CSSEstimator, ParCPDEstimator, ParCSSEstimator, SSE},
    models::{CatCPD, CatCPDS},
    types::{Set, States},
};

impl BE<'_, CatTable, usize> {
    // Fit a CPD given sufficient statistics.
    fn fit(
        states: &States,
        shape: &Array1<usize>,
        x: &Set<usize>,
        z: &Set<usize>,
        sample_statistics: CatCPDS,
        prior: usize,
    ) -> CatCPD {
        // Get the conditional counts.
        let n_xz = sample_statistics.sample_conditional_counts();
        // Marginalize the counts.
        let n_z = n_xz.sum_axis(Axis(1)).insert_axis(Axis(1));

        // Get the prior, as the alpha of the Dirichlet distribution.
        let alpha = prior;
        // Assert alpha is positive.
        assert!(alpha > 0, "Alpha must be positive.");

        // Cast alpha to floating point.
        let alpha = alpha as f64;

        // Add the prior to the counts.
        let n_xz = n_xz + alpha;
        let n_z = n_z + alpha * x.iter().map(|&i| shape[i]).product::<usize>() as f64;
        // Compute the parameters by normalizing the counts with the prior.
        let parameters = &n_xz / &n_z;

        // Compute the sample log-likelihood.
        let sample_log_likelihood = Some((&n_xz * parameters.ln()).sum());

        // Subset the conditioning labels, states and shape.
        let conditioning_states = z
            .iter()
            .map(|&i| {
                let (k, v) = states.get_index(i).unwrap();
                (k.clone(), v.clone())
            })
            .collect();
        // Get the labels of the conditioned variables.
        let states = x
            .iter()
            .map(|&i| {
                let (k, v) = states.get_index(i).unwrap();
                (k.clone(), v.clone())
            })
            .collect();

        // Wrap the sample statistics in an option.
        let sample_statistics = Some(sample_statistics);

        // Construct the CPD.
        CatCPD::with_optionals(
            states,
            conditioning_states,
            parameters,
            sample_statistics,
            sample_log_likelihood,
        )
    }
}

// Implement the CPD estimator for the BE struct.
macro_for!($type in [CatTable, CatIncTable, CatWtdTable] {

    impl CPDEstimator<CatCPD> for BE<'_, $type, ()> {
        #[inline]
        fn fit(&self, x: &Set<usize>, z: &Set<usize>) -> CatCPD {
            // Default to uniform prior.
            self.clone().with_prior(1).fit(x, z)
        }
    }

    impl CPDEstimator<CatCPD> for BE<'_, $type, usize> {
        #[inline]
        fn fit(&self, x: &Set<usize>, z: &Set<usize>) -> CatCPD {
            // Get (states, shape, prior).
            let (states, shape, prior) = (self.dataset.states(), self.dataset.shape(), self.prior);
            // Set sufficient statistics estimator.
            let sample_statistics = SSE::new(self.dataset);
            // Set missing handling method, if any.
            let sample_statistics = sample_statistics.with_missing_method(
                self.missing_method,
                self.missing_mechanism.clone()
            );
            // Compute sufficient statistics.
            let sample_statistics = sample_statistics.fit(x, z);
            // Fit the CPD given the sufficient statistics.
            BE::<'_, CatTable, _>::fit(states, shape, x, z, sample_statistics, prior)
        }
    }

    impl ParCPDEstimator<CatCPD> for BE<'_, $type, ()> {
        #[inline]
        fn par_fit(&self, x: &Set<usize>, z: &Set<usize>) -> CatCPD {
            // Default to uniform prior.
            self.clone().with_prior(1).fit(x, z)
        }
    }

    impl ParCPDEstimator<CatCPD> for BE<'_, $type, usize> {
        #[inline]
        fn par_fit(&self, x: &Set<usize>, z: &Set<usize>) -> CatCPD {
            // Get (states, shape, prior).
            let (states, shape, prior) = (self.dataset.states(), self.dataset.shape(), self.prior);
            // Set sufficient statistics estimator.
            let sample_statistics = SSE::new(self.dataset);
            // Set missing handling method, if any.
            let sample_statistics = sample_statistics.with_missing_method(
                self.missing_method,
                self.missing_mechanism.clone()
            );
            // Compute sufficient statistics in parallel.
            let sample_statistics = sample_statistics.par_fit(x, z);
            // Fit the CPD given the sufficient statistics.
            BE::<'_, CatTable, _>::fit(states, shape, x, z, sample_statistics, prior)
        }
    }

});