use crate::{
db::{
api::{mutation::AdditionOps, view::*},
graph::graph::Graph,
},
errors::GraphError,
prelude::NO_PROPS,
};
use rand::{rngs::StdRng, Rng, SeedableRng};
use raphtory_api::core::storage::timeindex::AsTime;
use raphtory_core::entities::GID;
pub fn erdos_renyi(nodes_to_add: usize, p: f64, seed: Option<u64>) -> Result<Graph, GraphError> {
let graph = Graph::new();
let mut rng;
if let Some(seed_value) = seed {
rng = StdRng::seed_from_u64(seed_value);
} else {
rng = StdRng::from_entropy();
}
let mut latest_time = graph.latest_time().map_or(0, |t| t.t());
for i in 0..nodes_to_add {
let id = GID::U64(i as u64);
latest_time += 1;
graph.add_node(latest_time, &id, NO_PROPS, None)?;
}
for i in 0..nodes_to_add {
let source_id = GID::U64(i as u64);
for j in (i + 1)..nodes_to_add {
let dst_id = GID::U64(j as u64);
let create_edge = rng.gen_bool(p);
if create_edge {
latest_time += 1;
graph.add_edge(latest_time, &source_id, &dst_id, NO_PROPS, None)?;
graph.add_edge(latest_time, &dst_id, &source_id, NO_PROPS, None)?;
}
}
}
Ok(graph)
}
#[cfg(test)]
mod tests {
use crate::{
db::api::view::{graph::GraphViewOps, node::NodeViewOps},
graphgen::erdos_renyi::erdos_renyi,
prelude::NodeStateOps,
};
#[test]
fn test_erdos_renyi_half_probability() {
let n_nodes = 20;
let p = 0.5;
let seed = Some(42);
let graph = erdos_renyi(n_nodes, p, seed).unwrap();
let node_count = graph.nodes().id().iter_values().count();
let edge_count = graph.edges().into_iter().count();
assert_eq!(node_count, n_nodes);
assert!(edge_count > 0);
assert!(edge_count <= n_nodes * (n_nodes - 1));
}
#[test]
fn test_erdos_renyi_zero_probability() {
let n_nodes = 20;
let p = 0.0;
let seed = Some(42);
let graph = erdos_renyi(n_nodes, p, seed).unwrap();
let edge_count = graph.edges().into_iter().count();
let node_count = graph.nodes().id().iter_values().count();
assert_eq!(node_count, n_nodes);
assert_eq!(edge_count, 0);
}
#[test]
fn test_erdos_renyi_full_probability() {
let n_nodes = 20;
let p = 1.0;
let seed = Some(42);
let graph = erdos_renyi(n_nodes, p, seed).unwrap();
let edge_count = graph.edges().into_iter().count();
let node_count = graph.nodes().id().iter_values().count();
assert_eq!(node_count, n_nodes);
assert_eq!(edge_count, (n_nodes * (n_nodes - 1)));
}
}