causal_hub/random/models/graphs/
directed.rs1use 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
10pub 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 pub fn new(rng: &'a mut R, labels: &'a Labels, probability: f64) -> Result<Self> {
31 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 let mut graph = DiGraph::empty(self.labels)?;
50
51 let n = self.labels.len();
53
54 (0..n)
56 .cartesian_product(0..n)
57 .filter(|(i, j)| i != j && self.rng.random_bool(self.probability))
59 .try_for_each(|(i, j)| -> Result<()> {
60 graph.add_edge(i, j)?;
62
63 Ok(())
64 })?;
65
66 Ok(graph)
67 }
68}
69
70pub 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 pub fn new(rng: &'a mut R, labels: &'a Labels, probability: f64) -> Result<Self> {
91 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 let mut graph = DiGraph::empty(self.labels)?;
110
111 let n = self.labels.len();
113
114 let mut order: Vec<_> = (0..n).collect();
116 order.shuffle(self.rng);
117
118 order
120 .into_iter()
121 .combinations(2)
122 .filter(|_| self.rng.random_bool(self.probability))
124 .try_for_each(|idx| -> Result<()> {
125 let (i, j) = (idx[0], idx[1]);
127 graph.add_edge(i, j)?;
128
129 Ok(())
130 })?;
131
132 Ok(graph)
133 }
134}