causal_hub/random/models/bayesian_network/categorical/
model.rs

1use rand::prelude::*;
2
3use crate::{
4    models::{BN, CatBN, CatSupport},
5    random::{Random, RngCatCPD, RngDag},
6    set,
7    types::{Error, Labels, Result},
8};
9
10/// A struct for random categorical Bayesian network generation.
11pub struct RngCatBN<'a, R>
12where
13    R: Rng,
14{
15    rng: &'a mut R,
16    support: &'a CatSupport,
17    alpha: f64,
18    probability: f64,
19}
20
21impl<'a, R> RngCatBN<'a, R>
22where
23    R: Rng,
24{
25    /// Creates a new `RngCatBN` instance.
26    ///
27    /// # Arguments
28    ///
29    /// * `rng` - A mutable reference to a random number generator.
30    /// * `support` - The support of the variables.
31    /// * `alpha` - The parameter of the Dirichlet distribution.
32    /// * `p` - The probability of generating an edge.
33    ///
34    /// # Errors
35    ///
36    /// * If `alpha` is not positive.
37    /// * If `p` is not in [0, 1].
38    ///
39    /// # Returns
40    ///
41    /// A new `RngCatBN` instance.
42    ///
43    pub fn new(
44        rng: &'a mut R,
45        support: &'a CatSupport,
46        alpha: f64,
47        probability: f64,
48    ) -> Result<Self> {
49        // Check if alpha is positive.
50        if alpha <= 0.0 {
51            return Err(Error::InvalidParameter("alpha", "must be positive"));
52        }
53        // Check if the probability is in [0, 1].
54        if !(0.0..=1.0).contains(&probability) {
55            return Err(Error::InvalidParameter("p", "must be in [0, 1]"));
56        }
57
58        Ok(Self {
59            rng,
60            support,
61            alpha,
62            probability,
63        })
64    }
65}
66
67impl<R> Random for RngCatBN<'_, R>
68where
69    R: Rng,
70{
71    type Output = Result<CatBN>;
72
73    fn random(&mut self) -> Self::Output {
74        // Get the labels of the variables.
75        let labels: Labels = self.support.keys().cloned().collect();
76
77        // Generate a random DAG.
78        let graph = RngDag::new(self.rng, &labels, self.probability)?.random()?;
79
80        // Generate the CPDs.
81        let cpds = labels
82            .iter()
83            .enumerate()
84            .map(|(i, x)| {
85                // Get the parents of the variable.
86                let pa_i = graph.parents(&set![i])?;
87                // Get the support of the variable.
88                let mut support = CatSupport::default();
89                support.insert(x.clone(), self.support[x].clone());
90                // Get the support of the conditioning variables.
91                let conditioning_support = pa_i
92                    .iter()
93                    .map(|&j| {
94                        let y = &labels[j];
95                        (y.clone(), self.support[y].clone())
96                    })
97                    .collect();
98                // Generate the random CPD.
99                RngCatCPD::new(self.rng, &support, &conditioning_support, self.alpha)?.random()
100            })
101            .collect::<Result<Vec<_>>>()?;
102
103        // Return the categorical Bayesian network.
104        CatBN::new(graph, cpds)
105    }
106}
causal_hub/random/models/bayesian_network/categorical/model.rs

causal_hub/random/models/bayesian_network/categorical/
model.rs
