Add barabasi_albert_graph random graph functions (#1007)

* Add barabasi_albert_graph random graph functions This commit adds new random graph functions to rustworkx and rustworkx-core to implement a random graph generator using the Barabási–Albert preferential attachment method. It takes an input graph (defaulting to a star graph) and then extends it to a given size. * Switch release note to use mpl_draw instead of graphviz_draw * Apply suggestions from code review Co-authored-by: Edwin Navarro <[email protected]> --------- Co-authored-by: Edwin Navarro <[email protected]>
Qiskit · Oct 18, 2023 · a424009 · a424009
1 parent f4ee4cf
commit a424009
Show file tree

Hide file tree

Showing 7 changed files with 366 additions and 2 deletions.
diff --git a/docs/source/api/random_graph_generator_functions.rst b/docs/source/api/random_graph_generator_functions.rst
@@ -11,3 +11,5 @@ Random Graph Generator Functions
     rustworkx.directed_gnm_random_graph
     rustworkx.undirected_gnm_random_graph
     rustworkx.random_geometric_graph
+    rustworkx.barabasi_albert_graph
+    rustworkx.directed_barabasi_albert_graph
diff --git a/releasenotes/notes/add-albert-graph-5a7d393e1fe18e9d.yaml b/releasenotes/notes/add-albert-graph-5a7d393e1fe18e9d.yaml
@@ -0,0 +1,20 @@
+---
+features:
+  - |
+    Added two new random graph generator functions,
+    :func:`.directed_barabasi_albert_graph` and :func:`.barabasi_albert_graph`,
+    to generate a random graph using Barabási–Albert preferential attachment to
+    extend an input graph. For example:
+
+        .. jupyter-execute::
+
+            import rustworkx
+            from rustworkx.visualization import mpl_draw
+
+            starting_graph = rustworkx.generators.path_graph(10)
+            random_graph = rustworkx.barabasi_albert_graph(20, 10, initial_graph=starting_graph)
+            mpl_draw(random_graph)
+  - |
+    Added a new function to the rustworkx-core module ``rustworkx_core::generators``
+    ``barabasi_albert_graph()`` which is used to generate a random graph
+    using Barabási–Albert preferential attachment to extend an input graph.
diff --git a/rustworkx-core/src/generators/mod.rs b/rustworkx-core/src/generators/mod.rs
@@ -56,6 +56,7 @@ pub use hexagonal_lattice_graph::hexagonal_lattice_graph;
 pub use lollipop_graph::lollipop_graph;
 pub use path_graph::path_graph;
 pub use petersen_graph::petersen_graph;
+pub use random_graph::barabasi_albert_graph;
 pub use random_graph::gnm_random_graph;
 pub use random_graph::gnp_random_graph;
 pub use random_graph::random_geometric_graph;

diff --git a/rustworkx-core/src/generators/random_graph.rs b/rustworkx-core/src/generators/random_graph.rs
@@ -15,12 +15,18 @@
 use std::hash::Hash;
 
 use petgraph::data::{Build, Create};
-use petgraph::visit::{Data, EdgeRef, GraphBase, GraphProp, IntoEdgeReferences, NodeIndexable};
+use petgraph::visit::{
+    Data, EdgeRef, GraphBase, GraphProp, IntoEdgeReferences, IntoEdgesDirected,
+    IntoNodeIdentifiers, NodeCount, NodeIndexable,
+};
+use petgraph::{Incoming, Outgoing};
 
+use hashbrown::HashSet;
 use rand::distributions::{Distribution, Uniform};
 use rand::prelude::*;
 use rand_pcg::Pcg64;
 
+use super::star_graph;
 use super::InvalidInputError;
 
 /// Generate a G<sub>np</sub> random graph, also known as an
@@ -395,10 +401,129 @@ where
     Ok(graph)
 }
 
+/// Generate a random Barabási–Albert preferential attachment algorithm
+///
+/// A graph of `n` nodes is grown by attaching new nodes each with `m`
+/// edges that are preferentially attached to existing nodes with high degree.
+/// If the graph is directed for the purposes of the extension algorithm all
+/// edges are treated as weak (meaning both incoming and outgoing).
+///
+/// The algorithm implemented in this function is described in:
+///
+/// A. L. Barabási and R. Albert "Emergence of scaling in random networks",
+/// Science 286, pp 509-512, 1999.
+///
+/// Arguments:
+///
+/// * `n` - The number of nodes to extend the graph to.
+/// * `m` - The number of edges to attach from a new node to existing nodes.
+/// * `seed` - An optional seed to use for the random number generator.
+/// * `initial_graph` - An optional starting graph to expand, if not specified
+///     a star graph of `m` nodes is generated and used. If specified the input
+///     graph is mutated by this function and is expected to be moved into this
+///     function.
+/// * `default_node_weight` - A callable that will return the weight to use
+///     for newly created nodes.
+/// * `default_edge_weight` - A callable that will return the weight object
+///     to use for newly created edges.
+///
+/// An `InvalidInput` error is returned under the following conditions. If `m < 1`
+/// or `m >= n` and if an `initial_graph` is specified and the number of nodes in
+/// `initial_graph` is `< m` or `> n`.
+///
+/// # Example
+/// ```rust
+/// use rustworkx_core::petgraph;
+/// use rustworkx_core::generators::barabasi_albert_graph;
+/// use rustworkx_core::generators::star_graph;
+///
+/// let graph: petgraph::graph::UnGraph<(), ()> = barabasi_albert_graph(
+///     20,
+///     12,
+///     Some(42),
+///     None,
+///     || {()},
+///     || {()},
+/// ).unwrap();
+/// assert_eq!(graph.node_count(), 20);
+/// assert_eq!(graph.edge_count(), 107);
+/// ```
+pub fn barabasi_albert_graph<G, T, F, H, M>(
+    n: usize,
+    m: usize,
+    seed: Option<u64>,
+    initial_graph: Option<G>,
+    mut default_node_weight: F,
+    mut default_edge_weight: H,
+) -> Result<G, InvalidInputError>
+where
+    G: Data<NodeWeight = T, EdgeWeight = M>
+        + NodeIndexable
+        + GraphProp
+        + NodeCount
+        + Build
+        + Create,
+    for<'b> &'b G: GraphBase<NodeId = G::NodeId> + IntoEdgesDirected + IntoNodeIdentifiers,
+    F: FnMut() -> T,
+    H: FnMut() -> M,
+    G::NodeId: Eq + Hash,
+{
+    if m < 1 || m >= n {
+        return Err(InvalidInputError {});
+    }
+    let mut rng: Pcg64 = match seed {
+        Some(seed) => Pcg64::seed_from_u64(seed),
+        None => Pcg64::from_entropy(),
+    };
+    let mut graph = match initial_graph {
+        Some(initial_graph) => initial_graph,
+        None => star_graph(
+            Some(m),
+            None,
+            &mut default_node_weight,
+            &mut default_edge_weight,
+            false,
+            false,
+        )?,
+    };
+    if graph.node_count() < m || graph.node_count() > n {
+        return Err(InvalidInputError {});
+    }
+
+    let mut repeated_nodes: Vec<G::NodeId> = graph
+        .node_identifiers()
+        .flat_map(|x| {
+            let degree = graph
+                .edges_directed(x, Outgoing)
+                .chain(graph.edges_directed(x, Incoming))
+                .count();
+            std::iter::repeat(x).take(degree)
+        })
+        .collect();
+    let mut source = graph.node_count();
+    while source < n {
+        let source_index = graph.add_node(default_node_weight());
+        let mut targets: HashSet<G::NodeId> = HashSet::with_capacity(m);
+        while targets.len() < m {
+            targets.insert(*repeated_nodes.choose(&mut rng).unwrap());
+        }
+        for target in &targets {
+            graph.add_edge(source_index, *target, default_edge_weight());
+        }
+        repeated_nodes.extend(targets);
+        repeated_nodes.extend(vec![source_index; m]);
+        source += 1
+    }
+    Ok(graph)
+}
+
 #[cfg(test)]
 mod tests {
     use crate::generators::InvalidInputError;
-    use crate::generators::{gnm_random_graph, gnp_random_graph, random_geometric_graph};
+    use crate::generators::{
+        barabasi_albert_graph, gnm_random_graph, gnp_random_graph, path_graph,
+        random_geometric_graph,
+    };
     use crate::petgraph;
 
     // Test gnp_random_graph
@@ -574,4 +699,54 @@ mod tests {
             Err(e) => assert_eq!(e, InvalidInputError),
         };
     }
+
+    #[test]
+    fn test_barabasi_albert_graph_starting_graph() {
+        let starting_graph: petgraph::graph::UnGraph<(), ()> =
+            path_graph(Some(40), None, || (), || (), false).unwrap();
+        let graph =
+            barabasi_albert_graph(500, 40, None, Some(starting_graph), || (), || ()).unwrap();
+        assert_eq!(graph.node_count(), 500);
+        assert_eq!(graph.edge_count(), 18439);
+    }
+
+    #[test]
+    fn test_barabasi_albert_graph_invalid_starting_size() {
+        match barabasi_albert_graph(
+            5,
+            40,
+            None,
+            None::<petgraph::graph::UnGraph<(), ()>>,
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        }
+    }
+
+    #[test]
+    fn test_barabasi_albert_graph_invalid_equal_starting_size() {
+        match barabasi_albert_graph(
+            5,
+            5,
+            None,
+            None::<petgraph::graph::UnGraph<(), ()>>,
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        }
+    }
+
+    #[test]
+    fn test_barabasi_albert_graph_invalid_starting_graph() {
+        let starting_graph: petgraph::graph::UnGraph<(), ()> =
+            path_graph(Some(4), None, || (), || (), false).unwrap();
+        match barabasi_albert_graph(500, 40, None, Some(starting_graph), || (), || ()) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        }
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -454,6 +454,8 @@ fn rustworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(directed_gnm_random_graph))?;
     m.add_wrapped(wrap_pyfunction!(undirected_gnm_random_graph))?;
     m.add_wrapped(wrap_pyfunction!(random_geometric_graph))?;
+    m.add_wrapped(wrap_pyfunction!(barabasi_albert_graph))?;
+    m.add_wrapped(wrap_pyfunction!(directed_barabasi_albert_graph))?;
     m.add_wrapped(wrap_pyfunction!(cycle_basis))?;
     m.add_wrapped(wrap_pyfunction!(simple_cycles))?;
     m.add_wrapped(wrap_pyfunction!(strongly_connected_components))?;

diff --git a/src/random_graph.rs b/src/random_graph.rs
@@ -379,3 +379,125 @@ pub fn random_geometric_graph(
     };
     Ok(graph)
 }
+
+/// Generate a random graph using Barabási–Albert preferential attachment
+///
+/// A graph is grown to $n$ nodes by adding new nodes each with $m$ edges that
+/// are preferentially attached to existing nodes with high degree. All the edges
+/// and nodes added to this graph will have weights of ``None``.
+///
+/// The algorithm performed by this function is described in:
+///
+/// A. L. Barabási and R. Albert "Emergence of scaling in random networks",
+/// Science 286, pp 509-512, 1999.
+///
+/// :param int n: The number of nodes to extend the graph to.
+/// :param int m: The number of edges to attach from a new node to existing nodes.
+/// :param int seed: An optional seed to use for the random number generator.
+/// :param PyGraph initial_graph: An optional initial graph to use as a starting
+///     point. :func:`.star_graph` is used to create an ``m`` node star graph
+///     to use as a starting point. If specified the input graph will be
+///     modified in place.
+///
+/// :return: A PyGraph object
+/// :rtype: PyGraph
+#[pyfunction]
+pub fn barabasi_albert_graph(
+    py: Python,
+    n: usize,
+    m: usize,
+    seed: Option<u64>,
+    initial_graph: Option<graph::PyGraph>,
+) -> PyResult<graph::PyGraph> {
+    let default_fn = || py.None();
+    if m < 1 {
+        return Err(PyValueError::new_err("m must be > 0"));
+    }
+    if m >= n {
+        return Err(PyValueError::new_err("m must be < n"));
+    }
+    let graph = match core_generators::barabasi_albert_graph(
+        n,
+        m,
+        seed,
+        initial_graph.map(|x| x.graph),
+        default_fn,
+        default_fn,
+    ) {
+        Ok(graph) => graph,
+        Err(_) => {
+            return Err(PyValueError::new_err(
+                "initial_graph has either less nodes than m, or more nodes than n",
+            ))
+        }
+    };
+    Ok(graph::PyGraph {
+        graph,
+        node_removed: false,
+        multigraph: true,
+        attrs: py.None(),
+    })
+}
+
+/// Generate a random graph using Barabási–Albert preferential attachment
+///
+/// A graph is grown to $n$ nodes by adding new nodes each with $m$ edges that
+/// are preferentially attached to existing nodes with high degree. All the edges
+/// and nodes added to this graph will have weights of ``None``. For the purposes
+/// of the extension algorithm all edges are treated as weak (meaning directionality
+/// isn't considered).
+///
+/// The algorithm performed by this function is described in:
+///
+/// A. L. Barabási and R. Albert "Emergence of scaling in random networks",
+/// Science 286, pp 509-512, 1999.
+///
+/// :param int n: The number of nodes to extend the graph to.
+/// :param int m: The number of edges to attach from a new node to existing nodes.
+/// :param int seed: An optional seed to use for the random number generator
+/// :param PyDiGraph initial_graph: An optional initial graph to use as a starting
+///     point. :func:`.star_graph` is used to create an ``m`` node star graph
+///     to use as a starting point. If specified the input graph will be
+///     modified in place.
+///
+/// :return: A PyDiGraph object
+/// :rtype: PyDiGraph
+#[pyfunction]
+pub fn directed_barabasi_albert_graph(
+    py: Python,
+    n: usize,
+    m: usize,
+    seed: Option<u64>,
+    initial_graph: Option<digraph::PyDiGraph>,
+) -> PyResult<digraph::PyDiGraph> {
+    let default_fn = || py.None();
+    if m < 1 {
+        return Err(PyValueError::new_err("m must be > 0"));
+    }
+    if m >= n {
+        return Err(PyValueError::new_err("m must be < n"));
+    }
+    let graph = match core_generators::barabasi_albert_graph(
+        n,
+        m,
+        seed,
+        initial_graph.map(|x| x.graph),
+        default_fn,
+        default_fn,
+    ) {
+        Ok(graph) => graph,
+        Err(_) => {
+            return Err(PyValueError::new_err(
+                "initial_graph has either less nodes than m, or more nodes than n",
+            ))
+        }
+    };
+    Ok(digraph::PyDiGraph {
+        graph,
+        node_removed: false,
+        check_cycle: false,
+        cycle_state: algo::DfsSpace::default(),
+        multigraph: false,
+        attrs: py.None(),
+    })
+}