Skip to main content

causal_hub/random/models/graphs/
directed.rs

1use itertools::Itertools;
2use rand::{Rng, RngExt, seq::SliceRandom};
3
4use crate::{
5    models::{DiGraph, Graph},
6    random::Random,
7    types::{Error, Labels, Result},
8};
9
10/// A struct representing a random directed graph generator.
11pub struct RngDiGraph<'a, R> {
12    rng: &'a mut R,
13    labels: &'a Labels,
14    probability: f64,
15}
16
17impl<'a, R> RngDiGraph<'a, R> {
18    /// Creates a new `RngDiGraph` 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 `RngDiGraph` 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 RngDiGraph<'_, R> {
45    type Output = Result<DiGraph>;
46
47    fn random(&mut self) -> Self::Output {
48        // Construct the empty graph.
49        let mut graph = DiGraph::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            .cartesian_product(0..n)
57            // ... filter out self-loops and sample edges ...
58            .filter(|(i, j)| i != j && self.rng.random_bool(self.probability))
59            .try_for_each(|(i, j)| -> Result<()> {
60                // ... add an edge.
61                graph.add_edge(i, j)?;
62
63                Ok(())
64            })?;
65
66        Ok(graph)
67    }
68}
69
70/// A struct representing a random directed acyclic graph generator.
71pub struct RngDag<'a, R> {
72    rng: &'a mut R,
73    labels: &'a Labels,
74    probability: f64,
75}
76
77impl<'a, R> RngDag<'a, R> {
78    /// Creates a new `RngDag` instance.
79    ///
80    /// # Arguments
81    ///
82    /// * `rng` - A mutable reference to a random number generator.
83    /// * `labels` - The labels of the graph.
84    /// * `p` - The probability of generating an edge.
85    ///
86    /// # Returns
87    ///
88    /// A new `RngDag` instance.
89    ///
90    pub fn new(rng: &'a mut R, labels: &'a Labels, probability: f64) -> Result<Self> {
91        // Check if the probability is in [0, 1].
92        if !(0.0..=1.0).contains(&probability) {
93            return Err(Error::InvalidParameter("p", "must be in [0, 1]"));
94        }
95
96        Ok(Self {
97            rng,
98            labels,
99            probability,
100        })
101    }
102}
103
104impl<R: Rng> Random for RngDag<'_, R> {
105    type Output = Result<DiGraph>;
106
107    fn random(&mut self) -> Self::Output {
108        // Construct the empty graph.
109        let mut graph = DiGraph::empty(self.labels)?;
110
111        // Get the number of vertices.
112        let n = self.labels.len();
113
114        // Generate a random topological order.
115        let mut order: Vec<_> = (0..n).collect();
116        order.shuffle(self.rng);
117
118        // For each pair of vertices in the topological order ...
119        order
120            .into_iter()
121            .combinations(2)
122            // ... sample edges ...
123            .filter(|_| self.rng.random_bool(self.probability))
124            .try_for_each(|idx| -> Result<()> {
125                // ... add an edge.
126                let (i, j) = (idx[0], idx[1]);
127                graph.add_edge(i, j)?;
128
129                Ok(())
130            })?;
131
132        Ok(graph)
133    }
134}