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

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

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