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

1use rand::prelude::*;
2
3use crate::{
4    models::{BN, CatSupport, MixedBN, MixedCPD, MixedSupport},
5    random::{Random, RngCatCPD, RngDag, RngGaussCPD},
6    set,
7    types::{Error, Labels, Map, Result},
8};
9
10/// A struct for random mixed Bayesian network generation.
11pub struct RngMixedBN<'a, R>
12where
13    R: Rng,
14{
15    rng: &'a mut R,
16    labels: &'a Labels,
17    support: &'a Map<String, MixedSupport>,
18    alpha: f64,
19    s_a: f64,
20    s_b: f64,
21    evidence: f64,
22    probability: f64,
23}
24
25impl<'a, R> RngMixedBN<'a, R>
26where
27    R: Rng,
28{
29    /// Creates a new `RngMixedBN` instance.
30    ///
31    /// # Arguments
32    ///
33    /// * `rng` - A mutable reference to a random number generator.
34    /// * `labels` - The labels of the variables.
35    /// * `support` - The support of the variables (`MixedSupport` per variable).
36    /// * `alpha` - The Dirichlet parameter for categorical CPDs (must be positive if any categorical).
37    /// * `s_a` - The standard deviation of regression coefficients for Gaussian CPDs.
38    /// * `s_b` - The standard deviation of the intercept for Gaussian CPDs.
39    /// * `e` - A small positive constant for covariance regularization.
40    /// * `p` - The probability of generating an edge.
41    ///
42    /// # Errors
43    ///
44    /// * If `alpha` is not positive.
45    /// * If `s_a`, `s_b`, or `e` are not positive.
46    /// * If `p` is not in [0, 1].
47    ///
48    /// # Returns
49    ///
50    /// A new `RngMixedBN` instance.
51    ///
52    #[allow(clippy::too_many_arguments)]
53    pub fn new(
54        rng: &'a mut R,
55        labels: &'a Labels,
56        support: &'a Map<String, MixedSupport>,
57        alpha: f64,
58        s_a: f64,
59        s_b: f64,
60        evidence: f64,
61        probability: f64,
62    ) -> Result<Self> {
63        if alpha <= 0.0 {
64            return Err(Error::InvalidParameter("alpha", "must be positive"));
65        }
66        if s_a <= 0.0 {
67            return Err(Error::InvalidParameter("s_a", "must be positive"));
68        }
69        if s_b <= 0.0 {
70            return Err(Error::InvalidParameter("s_b", "must be positive"));
71        }
72        if evidence <= 0.0 {
73            return Err(Error::InvalidParameter("e", "must be positive"));
74        }
75        if !(0.0..=1.0).contains(&probability) {
76            return Err(Error::InvalidParameter("p", "must be in [0, 1]"));
77        }
78
79        Ok(Self {
80            rng,
81            labels,
82            support,
83            alpha,
84            s_a,
85            s_b,
86            evidence,
87            probability,
88        })
89    }
90}
91
92impl<R> Random for RngMixedBN<'_, R>
93where
94    R: Rng,
95{
96    type Output = Result<MixedBN>;
97
98    fn random(&mut self) -> Self::Output {
99        let graph = RngDag::new(self.rng, self.labels, self.probability)?.random()?;
100
101        let cpds = self
102            .labels
103            .iter()
104            .enumerate()
105            .map(|(i, x)| {
106                let pa_i = graph.parents(&set![i])?;
107                let mixed_support = &self.support[x];
108
109                match mixed_support {
110                    MixedSupport::Categorical(cat_support) => {
111                        let mut support = CatSupport::default();
112                        support.insert(x.clone(), cat_support[x].clone());
113
114                        let conditioning_support: CatSupport = pa_i
115                            .iter()
116                            .map(|&j| {
117                                let y = &self.labels[j];
118                                match &self.support[y] {
119                                    MixedSupport::Categorical(stats) => {
120                                        (y.clone(), stats[y].clone())
121                                    }
122                                    _ => unreachable!("parents must match CPD type"),
123                                }
124                            })
125                            .collect();
126
127                        let distribution =
128                            RngCatCPD::new(self.rng, &support, &conditioning_support, self.alpha)?
129                                .random()?;
130                        Ok(MixedCPD::Categorical(distribution))
131                    }
132                    MixedSupport::Gaussian(_) => {
133                        let v_labels = crate::labels![x.clone()];
134                        let conditioning_labels: Labels =
135                            pa_i.iter().map(|&j| self.labels[j].clone()).collect();
136
137                        let distribution = RngGaussCPD::new(
138                            self.rng,
139                            &v_labels,
140                            &conditioning_labels,
141                            self.s_a,
142                            self.s_b,
143                            self.evidence,
144                        )?
145                        .random()?;
146                        Ok(MixedCPD::Gaussian(distribution))
147                    }
148                }
149            })
150            .collect::<Result<Vec<_>>>()?;
151
152        MixedBN::new(graph, cpds)
153    }
154}
causal_hub/random/models/bayesian_network/mixed/model.rs

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