causal_hub/random/datasets/table/categorical/
incomplete.rs

1use ndarray::prelude::*;
2use ndarray_stats::QuantileExt;
3use rand::{Rng, RngExt};
4use rand_distr::Uniform;
5
6use crate::{
7    datasets::{CatIncTable, CatTable, CatType, Dataset, IncDataset, MissingMechanism},
8    models::HasLabels,
9    random::Random,
10    types::{Error, Result},
11};
12
13/// A struct representing a random incomplete categorical table dataset generator.
14pub struct RngCatIncTable<'a, R> {
15    rng: &'a mut R,
16    dataset: &'a CatTable,
17    missing_mechanism: &'a MissingMechanism,
18    p_min: f64,
19    p_max: f64,
20}
21
22impl<'a, R: Rng> RngCatIncTable<'a, R> {
23    /// Creates a new `RngCatIncTable` instance.
24    ///
25    /// # Arguments
26    ///
27    /// * `rng` - A mutable reference to a random number generator.
28    /// * `dataset` - A reference to the complete categorical table dataset.
29    /// * `missing_mechanism` - A reference to the missingness mechanism.
30    /// * `p_min` - The minimum probability of missingness.
31    /// * `p_max` - The maximum probability of missingness.
32    ///
33    /// # Returns
34    ///
35    /// A new `RngCatIncTable` instance.
36    pub fn new(
37        rng: &'a mut R,
38        dataset: &'a CatTable,
39        missing_mechanism: &'a MissingMechanism,
40        p_min: f64,
41        p_max: f64,
42    ) -> Result<Self> {
43        // Check that dataset labels are equals to missing mechanism labels.
44        if dataset.labels() != missing_mechanism.labels() {
45            return Err(Error::InvalidParameter(
46                "missing_mechanism",
47                "labels do not match dataset labels",
48            ));
49        }
50        // Check that p_min and p_max are in [0, 1].
51        if !(0.0..=1.0).contains(&p_min) {
52            return Err(Error::InvalidParameter("p_min", "must be in [0, 1]"));
53        }
54        if !(0.0..=1.0).contains(&p_max) {
55            return Err(Error::InvalidParameter("p_max", "must be in [0, 1]"));
56        }
57        // Check that p_min is less than or equal to p_max.
58        if p_min > p_max {
59            return Err(Error::InvalidParameter(
60                "p_min",
61                "must be less than or equal to p_max",
62            ));
63        }
64
65        Ok(Self {
66            rng,
67            dataset,
68            missing_mechanism,
69            p_min,
70            p_max,
71        })
72    }
73}
74
75impl<R: Rng> Random for RngCatIncTable<'_, R> {
76    type Output = Result<CatIncTable>;
77
78    fn random(&mut self) -> Self::Output {
79        // Get the missing indicator.
80        const M: CatType = CatIncTable::MISSING;
81        // Get dataset support.
82        let support = self.dataset.support().clone();
83        // Get dataset values.
84        let mut values = self.dataset.values().clone();
85
86        // Define the uniform distributions for sampling.
87        let p_s = Uniform::new_inclusive(0., 1.)?;
88        let p_u = Uniform::new_inclusive(self.p_min, self.p_max)?;
89
90        // Iterate over the missing mechanism.
91        for (&x, pa_x) in self.missing_mechanism {
92            // Get mutable reference to the column X.
93            let mut c_x = values.column_mut(x);
94
95            // Check if the variable has no parents.
96            if pa_x.is_empty() {
97                // Sample a missingness probability.
98                let p_x = self.rng.sample(p_u);
99                // Modify the corresponding column.
100                c_x.iter_mut().for_each(|x| {
101                    if self.rng.sample(p_s) < p_x {
102                        *x = M;
103                    }
104                });
105                // Continue to the next variable.
106                continue;
107            }
108
109            // For each parent ...
110            for &z in pa_x {
111                // Get reference to the column Z.
112                let c_z = self.dataset.values().column(z);
113                // Get the mode of the parent via direct counting.
114                let mut m_z = Array::from_elem(CatType::MAX as usize, 0);
115                c_z.iter().for_each(|&z| m_z[z as usize] += 1);
116                let s_z = match m_z.argmax() {
117                    Ok(s_z) => s_z as CatType,
118                    // If the mode cannot be found, skip this parent.
119                    _ => continue,
120                };
121                // Modify the corresponding column.
122                azip!((x in &mut c_x, z in &c_z) {
123                    // Sample a missingness probability for X given parent Z.
124                    if
125                        (self.rng.sample(p_s) < self.p_max && *z == s_z) ||
126                        (self.rng.sample(p_s) < self.p_min && *z != s_z)
127                    {
128                        *x = M;
129                    }
130                });
131            }
132        }
133
134        // Return the incomplete dataset.
135        CatIncTable::new(support, values)
136    }
137}
causal_hub/random/datasets/table/categorical/incomplete.rs

causal_hub/random/datasets/table/categorical/
incomplete.rs
