Skip to main content

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

1use ndarray::prelude::*;
2use rand::{Rng, RngExt};
3use rand_distr::Gamma;
4
5use crate::{
6    models::{CatCPD, CatSupport},
7    random::Random,
8    types::{Error, Result},
9};
10
11/// A struct for random categorical CPD generation.
12pub struct RngCatCPD<'a, R>
13where
14    R: Rng,
15{
16    rng: &'a mut R,
17    support: &'a CatSupport,
18    conditioning_support: &'a CatSupport,
19    alpha: f64,
20}
21
22impl<'a, R> RngCatCPD<'a, R>
23where
24    R: Rng,
25{
26    /// Creates a new `RngCatCPD` instance.
27    ///
28    /// # Arguments
29    ///
30    /// * `rng` - A mutable reference to a random number generator.
31    /// * `support` - The support of the target variable.
32    /// * `conditioning_support` - The support of the conditioning variables.
33    /// * `alpha` - The parameter of the Dirichlet distribution.
34    ///
35    /// # Errors
36    ///
37    /// * If `alpha` is not positive.
38    ///
39    /// # Returns
40    ///
41    /// A new `RngCatCPD` instance.
42    ///
43    pub fn new(
44        rng: &'a mut R,
45        support: &'a CatSupport,
46        conditioning_support: &'a CatSupport,
47        alpha: 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
54        Ok(Self {
55            rng,
56            support,
57            conditioning_support,
58            alpha,
59        })
60    }
61}
62
63impl<R> Random for RngCatCPD<'_, R>
64where
65    R: Rng,
66{
67    type Output = Result<CatCPD>;
68
69    fn random(&mut self) -> Self::Output {
70        // Get the inner state sizes.
71        let model = self.support.values().map(|v| v.len()).product();
72        let n = self
73            .conditioning_support
74            .values()
75            .map(|v| v.len())
76            .product();
77
78        // Create the Gamma distribution.
79        let gamma = Gamma::new(self.alpha, 1.0)
80            .map_err(|evidence| Error::InvalidParameter("alpha", &evidence.to_string()))?;
81
82        // Sample the parameters.
83        let mut parameters = Array::from_shape_fn((n, model), |_| self.rng.sample(gamma));
84        // Normalize the parameters row-wise.
85        parameters /= &parameters.sum_axis(Axis(1)).insert_axis(Axis(1));
86
87        // Return the categorical CPD.
88        CatCPD::new(
89            self.support.clone(),
90            self.conditioning_support.clone(),
91            parameters,
92        )
93    }
94}