Skip to main content

causal_hub/random/models/graphs/
undirected.rs

1use itertools::Itertools;
2use rand::{Rng, RngExt};
3
4use crate::{
5    models::{Graph, UnGraph},
6    random::Random,
7    types::{Error, Labels, Result},
8};
9
10/// A struct representing a random undirected graph generator.
11pub struct RngUnGraph<'a, R> {
12    rng: &'a mut R,
13    labels: &'a Labels,
14    probability: f64,
15}
16
17impl<'a, R> RngUnGraph<'a, R> {
18    /// Creates a new `RngUnGraph` instance.
19    ///
20    /// # Arguments
21    ///
22    /// * `rng` - A mutable reference to a random number generator.
23    /// * `labels` - The labels of the graph.
24    /// * `p` - The probability of generating an edge.
25    ///
26    /// # Returns
27    ///
28    /// A new `RngUnGraph` instance.
29    ///
30    pub fn new(rng: &'a mut R, labels: &'a Labels, probability: f64) -> Result<Self> {
31        // Check if the probability is in [0, 1].
32        if !(0.0..=1.0).contains(&probability) {
33            return Err(Error::InvalidParameter("p", "must be in [0, 1]"));
34        }
35
36        Ok(Self {
37            rng,
38            labels,
39            probability,
40        })
41    }
42}
43
44impl<R: Rng> Random for RngUnGraph<'_, R> {
45    type Output = Result<UnGraph>;
46
47    fn random(&mut self) -> Self::Output {
48        // Construct the empty graph.
49        let mut graph = UnGraph::empty(self.labels)?;
50
51        // Get the number of vertices.
52        let n = self.labels.len();
53
54        // For each pair of vertices ...
55        (0..n)
56            .combinations(2)
57            // ... sample edges ...
58            .filter(|_| self.rng.random_bool(self.probability))
59            .try_for_each(|idx| -> Result<()> {
60                // ... add an edge.
61                let (i, j) = (idx[0], idx[1]);
62                graph.add_edge(i, j)?;
63
64                Ok(())
65            })?;
66
67        Ok(graph)
68    }
69}