From 189c9f639dc5cf76c136cbe1a0796e0e6e022e51 Mon Sep 17 00:00:00 2001
From: Simon Lizotte <simon.lizotte.1@ulaval.ca>
Date: Tue, 21 May 2024 12:48:54 -0400
Subject: [PATCH 1/5] wip generator first version

---
 rustworkx-core/src/generators/random_graph.rs | 163 ++++++++++++++++++
 1 file changed, 163 insertions(+)

diff --git a/rustworkx-core/src/generators/random_graph.rs b/rustworkx-core/src/generators/random_graph.rs
index b237d052d..8ab802d6c 100644
--- a/rustworkx-core/src/generators/random_graph.rs
+++ b/rustworkx-core/src/generators/random_graph.rs
@@ -619,6 +619,152 @@ where
     Ok(graph)
 }
 
+/// Generate a G<sub>np</sub> random graph, also known as an
+/// Erdős-Rényi graph or a binomial graph.
+///
+/// For number of nodes `n` and probability `p`, the G<sub>np</sub>
+/// graph algorithm creates `n` nodes, and for all the `n * (n - 1)` possible edges,
+/// each edge is created independently with probability `p`.
+/// In general, for any probability `p`, the expected number of edges returned
+/// is `m = p * n * (n - 1)`. If `p = 0` or `p = 1`, the returned
+/// graph is not random and will always be an empty or a complete graph respectively.
+/// An empty graph has zero edges and a complete directed graph has `n (n - 1)` edges.
+/// The run time is `O(n + m)` where `m` is the expected number of edges mentioned above.
+/// When `p = 0`, run time always reduces to `O(n)`, as the lower bound.
+/// When `p = 1`, run time always goes to `O(n + n * (n - 1))`, as the upper bound.
+///
+/// For `0 < p < 1`, the algorithm is based on the implementation of the networkx function
+/// ``fast_gnp_random_graph``,
+/// <https://github.com/networkx/networkx/blob/networkx-2.4/networkx/generators/random_graphs.py#L49-L120>
+///
+/// Vladimir Batagelj and Ulrik Brandes,
+///    "Efficient generation of large random networks",
+///    Phys. Rev. E, 71, 036113, 2005.
+///
+/// Arguments:
+///
+/// * `num_nodes` - The number of nodes for creating the random graph.
+/// * `probability` - The probability of creating an edge between two nodes as a float.
+/// * `seed` - An optional seed to use for the random number generator.
+/// * `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.
+///
+/// # Example
+/// ```rust
+/// use rustworkx_core::petgraph;
+/// use rustworkx_core::generators::gnp_random_graph;
+///
+/// let g: petgraph::graph::DiGraph<(), ()> = gnp_random_graph(
+///     20,
+///     1.0,
+///     None,
+///     || {()},
+///     || {()},
+/// ).unwrap();
+/// assert_eq!(g.node_count(), 20);
+/// assert_eq!(g.edge_count(), 20 * (20 - 1));
+/// ```
+pub fn sbm_random_graph<G, T, F, H, M>(
+    blocks: &[usize],
+    probabilities: &[Vec<f64>],
+    loops: bool,
+    seed: Option<u64>,
+    mut default_node_weight: F,
+    mut default_edge_weight: H,
+) -> Result<G, InvalidInputError>
+where
+    G: Build + Create + Data<NodeWeight = T, EdgeWeight = M> + NodeIndexable + GraphProp,
+    F: FnMut() -> T,
+    H: FnMut() -> M,
+    G::NodeId: Eq + Hash,
+{
+    let num_nodes = blocks.len();
+    if num_nodes == 0 {
+        return Err(InvalidInputError {});
+    }
+    let num_communities = probabilities.len();
+    if probabilities
+        .iter()
+        .any(|xs| xs.len() != num_communities || xs.iter().any(|&x| !(0. ..=1.).contains(&x)))
+    {
+        return Err(InvalidInputError {});
+    }
+    if blocks.iter().max().unwrap_or(&usize::MAX) >= &num_communities {
+        return Err(InvalidInputError {});
+    }
+    if blocks.len() != num_nodes {
+        return Err(InvalidInputError {});
+    }
+
+    let mut graph = G::with_capacity(num_nodes, num_nodes);
+    let directed = graph.is_directed();
+    if !directed && !symmetric_matrix(probabilities) {
+        return Err(InvalidInputError {});
+    }
+
+    for _ in 0..num_nodes {
+        graph.add_node(default_node_weight());
+    }
+    let mut rng: Pcg64 = match seed {
+        Some(seed) => Pcg64::seed_from_u64(seed),
+        None => Pcg64::from_entropy(),
+    };
+    let between = Uniform::new(0.0, 1.0);
+
+    let mut block_partition: Vec<Vec<usize>> = (0..num_communities).map(|_| Vec::new()).collect();
+    for (vertex, block) in blocks.iter().enumerate() {
+        block_partition[*block].push(vertex);
+    }
+
+    let block_pairs = block_partition
+        .iter()
+        .enumerate()
+        .zip(block_partition.iter().enumerate())
+        .filter(|((i, _), (j, _))| directed || i <= j);
+
+    for ((b1, block1), (b2, block2)) in block_pairs {
+        let prob = probabilities[b1][b2];
+        println!("block_len={x}", x=
+            block1
+                .iter()
+                .zip(block2.iter())
+                .filter(|(v, w)| v != w && directed || v < w || v == w && loops).count()
+        );
+        if prob > 0. {
+            block1
+                .iter()
+                .zip(block2.iter())
+                .filter(|(v, w)| v != w && directed || v < w || v == w && loops)
+                .filter(|_| between.sample(&mut rng) < prob)
+                .for_each(|(&v, &w)| {
+                    graph.add_edge(
+                        graph.from_index(v),
+                        graph.from_index(w),
+                        default_edge_weight(),
+                    );
+                });
+        }
+    }
+    Ok(graph)
+}
+
+fn symmetric_matrix<T: std::cmp::PartialEq>(mat: &[Vec<T>]) -> bool {
+    let n = mat.len();
+    for (i, row) in mat.iter().enumerate().take(n - 1) {
+        if row.len() != n {
+            return false;
+        }
+        for (j, m_ij) in row.iter().enumerate().skip(i + 1) {
+            if m_ij != &mat[j][i] {
+                return false;
+            }
+        }
+    }
+    true
+}
+
 #[cfg(test)]
 mod tests {
     use crate::generators::InvalidInputError;
@@ -628,6 +774,8 @@ mod tests {
     };
     use crate::petgraph;
 
+    use super::sbm_random_graph;
+
     // Test gnp_random_graph
 
     #[test]
@@ -916,4 +1064,19 @@ mod tests {
             Err(e) => assert_eq!(e, InvalidInputError),
         };
     }
+
+    // Test sbm_random_graph
+    #[test]
+    fn test_sbm_complete_block() {
+        let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![1., 0.], vec![0., 1.]],
+            true,
+            None,
+            || (),
+            || (),
+        ).unwrap();
+        assert_eq!(g.node_count(), 3);
+        assert_eq!(g.edge_count(), 5);
+    }
 }

From 543c0b934113a91224ee2a63f7aa118b7afe5f1a Mon Sep 17 00:00:00 2001
From: Simon Lizotte <simon.lizotte.1@ulaval.ca>
Date: Wed, 22 May 2024 13:50:40 -0400
Subject: [PATCH 2/5] add sbm generator to rustworkx and rustworkx-core

---
 .../api/random_graph_generator_functions.rst  |   2 +
 .../sbm-random-graph-bf7ccd8e938f4218.yaml    |   9 +
 rustworkx-core/src/generators/mod.rs          |   1 +
 rustworkx-core/src/generators/random_graph.rs | 279 +++++++++++++-----
 rustworkx/__init__.pyi                        |   2 +
 rustworkx/rustworkx.pyi                       |  14 +
 src/lib.rs                                    |   2 +
 src/random_graph.rs                           | 108 +++++++
 tests/test_random.py                          |  65 ++++
 9 files changed, 414 insertions(+), 68 deletions(-)
 create mode 100644 releasenotes/notes/sbm-random-graph-bf7ccd8e938f4218.yaml

diff --git a/docs/source/api/random_graph_generator_functions.rst b/docs/source/api/random_graph_generator_functions.rst
index ba6823cfd..f02665d84 100644
--- a/docs/source/api/random_graph_generator_functions.rst
+++ b/docs/source/api/random_graph_generator_functions.rst
@@ -10,6 +10,8 @@ Random Graph Generator Functions
     rustworkx.undirected_gnp_random_graph
     rustworkx.directed_gnm_random_graph
     rustworkx.undirected_gnm_random_graph
+    rustworkx.directed_sbm_random_graph
+    rustworkx.undirected_sbm_random_graph
     rustworkx.random_geometric_graph
     rustworkx.barabasi_albert_graph
     rustworkx.directed_barabasi_albert_graph
diff --git a/releasenotes/notes/sbm-random-graph-bf7ccd8e938f4218.yaml b/releasenotes/notes/sbm-random-graph-bf7ccd8e938f4218.yaml
new file mode 100644
index 000000000..8ec9490a4
--- /dev/null
+++ b/releasenotes/notes/sbm-random-graph-bf7ccd8e938f4218.yaml
@@ -0,0 +1,9 @@
+features:
+  - |
+    Adds new random graph generator in rustworkx for the stochastic block model.
+    There is a generator for directed :func:`.directed_sbm_random_graph` and
+    undirected graphs :func:`.undirected_sbm_random_graph`.
+  - |
+    Adds new function ``sbm_random_graph`` to the rustworkx-core module
+    ``rustworkx_core::generators`` that samples a graph from the stochastic
+    block model.
diff --git a/rustworkx-core/src/generators/mod.rs b/rustworkx-core/src/generators/mod.rs
index d0a2c8703..164133b5e 100644
--- a/rustworkx-core/src/generators/mod.rs
+++ b/rustworkx-core/src/generators/mod.rs
@@ -61,4 +61,5 @@ pub use random_graph::gnm_random_graph;
 pub use random_graph::gnp_random_graph;
 pub use random_graph::random_bipartite_graph;
 pub use random_graph::random_geometric_graph;
+pub use random_graph::sbm_random_graph;
 pub use star_graph::star_graph;
diff --git a/rustworkx-core/src/generators/random_graph.rs b/rustworkx-core/src/generators/random_graph.rs
index 8ab802d6c..a140a0345 100644
--- a/rustworkx-core/src/generators/random_graph.rs
+++ b/rustworkx-core/src/generators/random_graph.rs
@@ -619,32 +619,20 @@ where
     Ok(graph)
 }
 
-/// Generate a G<sub>np</sub> random graph, also known as an
-/// Erdős-Rényi graph or a binomial graph.
-///
-/// For number of nodes `n` and probability `p`, the G<sub>np</sub>
-/// graph algorithm creates `n` nodes, and for all the `n * (n - 1)` possible edges,
-/// each edge is created independently with probability `p`.
-/// In general, for any probability `p`, the expected number of edges returned
-/// is `m = p * n * (n - 1)`. If `p = 0` or `p = 1`, the returned
-/// graph is not random and will always be an empty or a complete graph respectively.
-/// An empty graph has zero edges and a complete directed graph has `n (n - 1)` edges.
-/// The run time is `O(n + m)` where `m` is the expected number of edges mentioned above.
-/// When `p = 0`, run time always reduces to `O(n)`, as the lower bound.
-/// When `p = 1`, run time always goes to `O(n + n * (n - 1))`, as the upper bound.
-///
-/// For `0 < p < 1`, the algorithm is based on the implementation of the networkx function
-/// ``fast_gnp_random_graph``,
-/// <https://github.com/networkx/networkx/blob/networkx-2.4/networkx/generators/random_graphs.py#L49-L120>
+/// Generate a graph from the stochastic block model.
 ///
-/// Vladimir Batagelj and Ulrik Brandes,
-///    "Efficient generation of large random networks",
-///    Phys. Rev. E, 71, 036113, 2005.
+/// The stochastic block model is a generalization of the G<sub>np</sub> graph model
+/// (see [rustworkx_core::generators::gnp_random_graph] ). The connection probability of
+/// nodes `u` and `v` depends on their block (or community) and is given by
+/// `probabilities[blocks[u]][blocks[v]]`. The number of nodes and the number of blocks
+/// are inferred from `blocks`.
 ///
 /// Arguments:
 ///
-/// * `num_nodes` - The number of nodes for creating the random graph.
-/// * `probability` - The probability of creating an edge between two nodes as a float.
+/// * `blocks` - Block membership (between 0 and B-1) of each node.
+/// * `probabilities` - B x B matrix that contains the connection probability between
+///     nodes of different blocks. Must be symmetric for undirected graphs.
+/// * `loops` - Determines whether the graph can have loops or not.
 /// * `seed` - An optional seed to use for the random number generator.
 /// * `default_node_weight` - A callable that will return the weight to use
 ///     for newly created nodes.
@@ -654,17 +642,19 @@ where
 /// # Example
 /// ```rust
 /// use rustworkx_core::petgraph;
-/// use rustworkx_core::generators::gnp_random_graph;
+/// use rustworkx_core::generators::sbm_random_graph;
 ///
-/// let g: petgraph::graph::DiGraph<(), ()> = gnp_random_graph(
-///     20,
-///     1.0,
-///     None,
-///     || {()},
-///     || {()},
-/// ).unwrap();
-/// assert_eq!(g.node_count(), 20);
-/// assert_eq!(g.edge_count(), 20 * (20 - 1));
+/// let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+///     &vec![1, 0, 1],
+///     &[vec![0., 1.], vec![0., 1.]],
+///     true,
+///     Some(10),
+///     || (),
+///     || (),
+/// )
+/// .unwrap();
+/// assert_eq!(g.node_count(), 3);
+/// assert_eq!(g.edge_count(), 6);
 /// ```
 pub fn sbm_random_graph<G, T, F, H, M>(
     blocks: &[usize],
@@ -714,37 +704,28 @@ where
     let between = Uniform::new(0.0, 1.0);
 
     let mut block_partition: Vec<Vec<usize>> = (0..num_communities).map(|_| Vec::new()).collect();
-    for (vertex, block) in blocks.iter().enumerate() {
-        block_partition[*block].push(vertex);
+    for (node, block) in blocks.iter().enumerate() {
+        block_partition[*block].push(node);
     }
 
-    let block_pairs = block_partition
-        .iter()
-        .enumerate()
-        .zip(block_partition.iter().enumerate())
-        .filter(|((i, _), (j, _))| directed || i <= j);
-
-    for ((b1, block1), (b2, block2)) in block_pairs {
-        let prob = probabilities[b1][b2];
-        println!("block_len={x}", x=
-            block1
-                .iter()
-                .zip(block2.iter())
-                .filter(|(v, w)| v != w && directed || v < w || v == w && loops).count()
-        );
-        if prob > 0. {
-            block1
-                .iter()
-                .zip(block2.iter())
-                .filter(|(v, w)| v != w && directed || v < w || v == w && loops)
-                .filter(|_| between.sample(&mut rng) < prob)
-                .for_each(|(&v, &w)| {
-                    graph.add_edge(
-                        graph.from_index(v),
-                        graph.from_index(w),
-                        default_edge_weight(),
-                    );
-                });
+    for (v, &b_v) in blocks.iter().enumerate().take(if directed || loops {
+        num_nodes
+    } else {
+        num_nodes - 1
+    }) {
+        for (w, &b_w) in blocks
+            .iter()
+            .enumerate()
+            .skip(if directed { 0 } else { v })
+            .filter(|&(w, _)| w != v || loops)
+        {
+            if between.sample(&mut rng) < probabilities[b_v][b_w] {
+                graph.add_edge(
+                    graph.from_index(v),
+                    graph.from_index(w),
+                    default_edge_weight(),
+                );
+            }
         }
     }
     Ok(graph)
@@ -770,12 +751,10 @@ mod tests {
     use crate::generators::InvalidInputError;
     use crate::generators::{
         barabasi_albert_graph, gnm_random_graph, gnp_random_graph, path_graph,
-        random_bipartite_graph, random_geometric_graph,
+        random_bipartite_graph, random_geometric_graph, sbm_random_graph,
     };
     use crate::petgraph;
 
-    use super::sbm_random_graph;
-
     // Test gnp_random_graph
 
     #[test]
@@ -1067,16 +1046,180 @@ mod tests {
 
     // Test sbm_random_graph
     #[test]
-    fn test_sbm_complete_block() {
+    fn test_sbm_directed_complete_blocks() {
         let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
             &vec![1, 0, 1],
-            &[vec![1., 0.], vec![0., 1.]],
+            &[vec![0., 1.], vec![0., 1.]],
             true,
-            None,
+            Some(10),
+            || (),
+            || (),
+        )
+        .unwrap();
+        assert_eq!(g.node_count(), 3);
+        assert_eq!(g.edge_count(), 6);
+        for (u, v) in [(0, 0), (0, 2), (2, 0), (2, 2), (1, 0), (1, 2)] {
+            assert_eq!(g.contains_edge(u.into(), v.into()), true);
+        }
+        assert_eq!(g.contains_edge(0.into(), 1.into()), false);
+        assert_eq!(g.contains_edge(2.into(), 1.into()), false);
+    }
+
+    #[test]
+    fn test_sbm_directed_complete_blocks_loops() {
+        let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![0., 1.], vec![0., 1.]],
+            true,
+            Some(10),
             || (),
             || (),
-        ).unwrap();
+        )
+        .unwrap();
+        assert_eq!(g.node_count(), 3);
+        assert_eq!(g.edge_count(), 6);
+        for (u, v) in [(0, 0), (0, 2), (2, 0), (2, 2), (1, 0), (1, 2)] {
+            assert_eq!(g.contains_edge(u.into(), v.into()), true);
+        }
+        assert_eq!(g.contains_edge(0.into(), 1.into()), false);
+        assert_eq!(g.contains_edge(2.into(), 1.into()), false);
+    }
+
+    #[test]
+    fn test_sbm_undirected_complete_blocks_loops() {
+        let g = sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![0., 1.], vec![1., 1.]],
+            true,
+            Some(10),
+            || (),
+            || (),
+        )
+        .unwrap();
         assert_eq!(g.node_count(), 3);
         assert_eq!(g.edge_count(), 5);
+        for (u, v) in [(0, 0), (0, 2), (2, 2), (1, 0), (1, 2)] {
+            assert_eq!(g.contains_edge(u.into(), v.into()), true);
+        }
+        assert_eq!(g.contains_edge(1.into(), 1.into()), false);
+    }
+
+    #[test]
+    fn test_sbm_directed_complete_blocks_noloops() {
+        let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![0., 1.], vec![0., 1.]],
+            false,
+            Some(10),
+            || (),
+            || (),
+        )
+        .unwrap();
+        assert_eq!(g.node_count(), 3);
+        assert_eq!(g.edge_count(), 4);
+        for (u, v) in [(0, 2), (2, 0), (1, 0), (1, 2)] {
+            assert_eq!(g.contains_edge(u.into(), v.into()), true);
+        }
+        assert_eq!(g.contains_edge(0.into(), 1.into()), false);
+        assert_eq!(g.contains_edge(2.into(), 1.into()), false);
+        for u in 0..2 {
+            assert_eq!(g.contains_edge(u.into(), u.into()), false);
+        }
+    }
+
+    #[test]
+    fn test_sbm_undirected_complete_blocks_noloops() {
+        let g = sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![0., 1.], vec![1., 1.]],
+            false,
+            Some(10),
+            || (),
+            || (),
+        )
+        .unwrap();
+        assert_eq!(g.node_count(), 3);
+        assert_eq!(g.edge_count(), 3);
+        for (u, v) in [(0, 2), (1, 0), (1, 2)] {
+            assert_eq!(g.contains_edge(u.into(), v.into()), true);
+        }
+        for u in 0..2 {
+            assert_eq!(g.contains_edge(u.into(), u.into()), false);
+        }
+    }
+
+    #[test]
+    fn test_sbm_block_outofrange_error() {
+        match sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 2],
+            &[vec![0., 1.], vec![1., 1.]],
+            true,
+            Some(10),
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
+    }
+
+    #[test]
+    fn test_sbm_invalid_matrix_error() {
+        match sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![0., 1.], vec![1.]],
+            true,
+            Some(10),
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
+    }
+
+    #[test]
+    fn test_sbm_asymmetric_matrix_error() {
+        match sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![0., 1.], vec![0., 1.]],
+            true,
+            Some(10),
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
+    }
+
+    #[test]
+    fn test_sbm_invalid_probability_error() {
+        match sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
+            &vec![1, 0, 1],
+            &[vec![0., 1.], vec![0., -1.]],
+            true,
+            Some(10),
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
+    }
+
+    #[test]
+    fn test_sbm_empty_error() {
+        match sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+            &vec![],
+            &[],
+            true,
+            Some(10),
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
     }
 }
diff --git a/rustworkx/__init__.pyi b/rustworkx/__init__.pyi
index f5bcdeb71..6db3aadc3 100644
--- a/rustworkx/__init__.pyi
+++ b/rustworkx/__init__.pyi
@@ -127,6 +127,8 @@ from .rustworkx import directed_gnm_random_graph as directed_gnm_random_graph
 from .rustworkx import undirected_gnm_random_graph as undirected_gnm_random_graph
 from .rustworkx import directed_gnp_random_graph as directed_gnp_random_graph
 from .rustworkx import undirected_gnp_random_graph as undirected_gnp_random_graph
+from .rustworkx import directed_sbm_random_graph as directed_sbm_random_graph
+from .rustworkx import undirected_sbm_random_graph as undirected_sbm_random_graph
 from .rustworkx import random_geometric_graph as random_geometric_graph
 from .rustworkx import barabasi_albert_graph as barabasi_albert_graph
 from .rustworkx import directed_barabasi_albert_graph as directed_barabasi_albert_graph
diff --git a/rustworkx/rustworkx.pyi b/rustworkx/rustworkx.pyi
index 2edcdde67..8916dba11 100644
--- a/rustworkx/rustworkx.pyi
+++ b/rustworkx/rustworkx.pyi
@@ -549,6 +549,20 @@ def undirected_gnp_random_graph(
     /,
     seed: int | None = ...,
 ) -> PyGraph: ...
+def directed_sbm_random_graph(
+    blocks: list[int],
+    probabilities: list[list[float]],
+    loops: bool,
+    /,
+    seed: int | None = ...,
+) -> PyDiGraph: ...
+def undirected_sbm_random_graph(
+    blocks: list[int],
+    probabilities: list[list[float]],
+    loops: bool,
+    /,
+    seed: int | None = ...,
+) -> PyGraph: ...
 def random_geometric_graph(
     num_nodes: int,
     radius: float,
diff --git a/src/lib.rs b/src/lib.rs
index 096f30072..17c3a17a6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -476,6 +476,8 @@ fn rustworkx(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(undirected_gnp_random_graph))?;
     m.add_wrapped(wrap_pyfunction!(directed_gnm_random_graph))?;
     m.add_wrapped(wrap_pyfunction!(undirected_gnm_random_graph))?;
+    m.add_wrapped(wrap_pyfunction!(undirected_sbm_random_graph))?;
+    m.add_wrapped(wrap_pyfunction!(directed_sbm_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))?;
diff --git a/src/random_graph.rs b/src/random_graph.rs
index 39520a23e..b1cff10a6 100644
--- a/src/random_graph.rs
+++ b/src/random_graph.rs
@@ -273,6 +273,114 @@ pub fn undirected_gnm_random_graph(
     })
 }
 
+/// Return a directed graph from the stochastic block model.
+///
+/// The stochastic block model is a generalization of the G<sub>np</sub> graph model
+/// (see [rustworkx.undirected_gnp_random_graph] ). The connection probability of
+/// nodes `u` and `v` depends on their block (or community) and is given by
+/// `probabilities[blocks[u]][blocks[v]]`. The number of nodes and the number of blocks
+/// are inferred from `blocks`.
+///
+/// This algorithm has a time complexity of :math:`O(n^2)` for :math:`n` nodes.
+///
+/// Arguments:
+///
+/// :param list[int] blocks: Block membership (between 0 and B-1) of each node.
+/// :param list[list[float]] probabilities: Matrix B x B that contains the
+///     connection probability between nodes of different blocks.
+/// :param bool loops: Determines whether the graph can have loops or not.
+/// :param int seed:  An optional seed to use for the random number generator.
+///
+/// :return: A PyDiGraph object
+/// :rtype: PyDiGraph
+#[pyfunction]
+#[pyo3(text_signature = "(blocks, probabilities, loops, /, seed=None)")]
+pub fn directed_sbm_random_graph(
+    py: Python,
+    blocks: Vec<usize>,
+    probabilities: Vec<Vec<f64>>,
+    loops: bool,
+    seed: Option<u64>,
+) -> PyResult<digraph::PyDiGraph> {
+    let default_fn = || py.None();
+    let graph: StablePyGraph<Directed> = match core_generators::sbm_random_graph(
+        &blocks,
+        &probabilities,
+        loops,
+        seed,
+        default_fn,
+        default_fn,
+    ) {
+        Ok(graph) => graph,
+        Err(_) => {
+            return Err(PyValueError::new_err(
+                "invalid blocks or probabilities input",
+            ))
+        }
+    };
+    Ok(digraph::PyDiGraph {
+        graph,
+        node_removed: false,
+        check_cycle: false,
+        cycle_state: algo::DfsSpace::default(),
+        multigraph: false,
+        attrs: py.None(),
+    })
+}
+
+/// Return an undirected graph from the stochastic block model.
+///
+/// The stochastic block model is a generalization of the G<sub>np</sub> graph model
+/// (see [rustworkx.undirected_gnp_random_graph] ). The connection probability of
+/// nodes `u` and `v` depends on their block (or community) and is given by
+/// `probabilities[blocks[u]][blocks[v]]`. The number of nodes and the number of blocks
+/// are inferred from `blocks`.
+///
+/// This algorithm has a time complexity of :math:`O(n^2)` for :math:`n` nodes.
+///
+/// Arguments:
+///
+/// :param list[int] blocks: Block membership (between 0 and B-1) of each node.
+/// :param list[list[float]] probabilities: Symmetric B x B matrix that contains
+///     the connection probability between nodes of different blocks.
+/// :param bool loops: Determines whether the graph can have loops or not.
+/// :param int seed:  An optional seed to use for the random number generator.
+///
+/// :return: A PyGraph object
+/// :rtype: PyGraph
+#[pyfunction]
+#[pyo3(text_signature = "(blocks, probabilities, loops, /, seed=None)")]
+pub fn undirected_sbm_random_graph(
+    py: Python,
+    blocks: Vec<usize>,
+    probabilities: Vec<Vec<f64>>,
+    loops: bool,
+    seed: Option<u64>,
+) -> PyResult<graph::PyGraph> {
+    let default_fn = || py.None();
+    let graph: StablePyGraph<Undirected> = match core_generators::sbm_random_graph(
+        &blocks,
+        &probabilities,
+        loops,
+        seed,
+        default_fn,
+        default_fn,
+    ) {
+        Ok(graph) => graph,
+        Err(_) => {
+            return Err(PyValueError::new_err(
+                "invalid blocks or probabilities input",
+            ))
+        }
+    };
+    Ok(graph::PyGraph {
+        graph,
+        node_removed: false,
+        multigraph: false,
+        attrs: py.None(),
+    })
+}
+
 #[inline]
 fn pnorm(x: f64, p: f64) -> f64 {
     if p == 1.0 || p == std::f64::INFINITY {
diff --git a/tests/test_random.py b/tests/test_random.py
index 601ccc9d3..2f821f082 100644
--- a/tests/test_random.py
+++ b/tests/test_random.py
@@ -176,6 +176,71 @@ def test_random_gnm_undirected_payload(self):
         self.assertEqual(graph.nodes(), [0, 1, 2])
 
 
+class TestRandomSBM(unittest.TestCase):
+    def test_undirected_sbm_complete_blocks_loops(self):
+        graph = rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 1], [1, 0]], True)
+        self.assertEqual(len(graph), 3)
+        self.assertEqual(len(graph.edges()), 5)
+        for i in range(2):
+            for j in range(i, 2):
+                if (i, j) != (1, 1):
+                    self.assertTrue(graph.has_edge(i, j))
+        self.assertFalse(graph.has_edge(1, 1))
+
+    def test_directed_sbm_complete_blocks_loops(self):
+        graph = rustworkx.directed_sbm_random_graph([0, 1, 0], [[0, 0], [1, 1]], True)
+        self.assertEqual(len(graph), 3)
+        self.assertEqual(len(graph.edges()), 3)
+        self.assertEqual(set(graph.edge_list()), set([(1, 1), (1, 0), (1, 2)]))
+
+    def test_undirected_sbm_complete_blocks_noloops(self):
+        graph = rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 1], [1, 0]], False)
+        self.assertEqual(len(graph), 3)
+        self.assertEqual(len(graph.edges()), 3)
+        for i in range(2):
+            for j in range(i, 2):
+                if i != j:
+                    self.assertTrue(graph.has_edge(i, j))
+
+    def test_directed_sbm_complete_blocks_noloops(self):
+        graph = rustworkx.directed_sbm_random_graph([0, 1, 0], [[0, 0], [1, 1]], False)
+        self.assertEqual(len(graph), 3)
+        self.assertEqual(len(graph.edges()), 2)
+        self.assertEqual(set(graph.edge_list()), set([(1, 0), (1, 2)]))
+
+    def test_undirected_sbm_asymmetric_probabilities_error(self):
+        with self.assertRaises(ValueError):
+            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[0, 0], [1, 1]], True)
+
+    def test_sbm_out_of_range_blocks_error(self):
+        with self.assertRaises(ValueError):
+            rustworkx.undirected_sbm_random_graph([0, 2, 0], [[1, 0], [0, 1]], True)
+        with self.assertRaises(ValueError):
+            rustworkx.directed_sbm_random_graph([0, 2, 0], [[1, 0], [0, 1]], True)
+
+    def test_sbm_invalid_matrix(self):
+        with self.assertRaises(ValueError):
+            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1, 0]], True)
+        with self.assertRaises(ValueError):
+            rustworkx.directed_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1, 0]], True)
+
+    def test_sbm_invalid_probabilities(self):
+        with self.assertRaises(ValueError):
+            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1.5]], True)
+        with self.assertRaises(ValueError):
+            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[-1, 0], [0, 1]], True)
+        with self.assertRaises(ValueError):
+            rustworkx.directed_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1.5]], True)
+        with self.assertRaises(ValueError):
+            rustworkx.directed_sbm_random_graph([0, 1, 0], [[-1, 0], [0, 1]], True)
+
+    def test_sbm_empty(self):
+        with self.assertRaises(ValueError):
+            rustworkx.undirected_sbm_random_graph([], [], True)
+        with self.assertRaises(ValueError):
+            rustworkx.directed_sbm_random_graph([], [], True)
+
+
 class TestGeometricRandomGraph(unittest.TestCase):
     def test_random_geometric_empty(self):
         graph = rustworkx.random_geometric_graph(20, 0)

From 53a6ab6e79bbd30c1eefda86b0592278cc3cd5fd Mon Sep 17 00:00:00 2001
From: Simon Lizotte <simon.lizotte.1@ulaval.ca>
Date: Wed, 22 May 2024 14:35:38 -0400
Subject: [PATCH 3/5] fix formatting error

---
 src/random_graph.rs | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/random_graph.rs b/src/random_graph.rs
index 455dd34d2..ce8857fd0 100644
--- a/src/random_graph.rs
+++ b/src/random_graph.rs
@@ -275,18 +275,18 @@ pub fn undirected_gnm_random_graph(
 
 /// Return a directed graph from the stochastic block model.
 ///
-/// The stochastic block model is a generalization of the G<sub>np</sub> graph model
+/// The stochastic block model is a generalization of the :math:`G(n,p)` graph model
 /// (see [rustworkx.undirected_gnp_random_graph] ). The connection probability of
-/// nodes `u` and `v` depends on their block (or community) and is given by
-/// `probabilities[blocks[u]][blocks[v]]`. The number of nodes and the number of blocks
-/// are inferred from `blocks`.
+/// nodes ``u`` and ``v`` depends on their block (or community) and is given by
+/// ``probabilities[blocks[u]][blocks[v]]``. The number of nodes and the number of
+/// blocks are inferred from ``blocks``.
 ///
 /// This algorithm has a time complexity of :math:`O(n^2)` for :math:`n` nodes.
 ///
 /// Arguments:
 ///
 /// :param list[int] blocks: Block membership (between 0 and B-1) of each node.
-/// :param list[list[float]] probabilities: Matrix B x B that contains the
+/// :param list[list[float]] probabilities: B x B matrix that contains the
 ///     connection probability between nodes of different blocks.
 /// :param bool loops: Determines whether the graph can have loops or not.
 /// :param int seed:  An optional seed to use for the random number generator.
@@ -330,11 +330,11 @@ pub fn directed_sbm_random_graph(
 
 /// Return an undirected graph from the stochastic block model.
 ///
-/// The stochastic block model is a generalization of the G<sub>np</sub> graph model
+/// The stochastic block model is a generalization of the :math:`G(n,p)` graph model
 /// (see [rustworkx.undirected_gnp_random_graph] ). The connection probability of
-/// nodes `u` and `v` depends on their block (or community) and is given by
-/// `probabilities[blocks[u]][blocks[v]]`. The number of nodes and the number of blocks
-/// are inferred from `blocks`.
+/// nodes ``u`` and ``v`` depends on their block (or community) and is given by
+/// ``probabilities[blocks[u]][blocks[v]]``. The number of nodes and the number of
+/// blocks are inferred from ``blocks``.
 ///
 /// This algorithm has a time complexity of :math:`O(n^2)` for :math:`n` nodes.
 ///

From a8bef7986b96183f1947531ca1f04a112c7ed5de Mon Sep 17 00:00:00 2001
From: Simon Lizotte <simon.lizotte.1@ulaval.ca>
Date: Fri, 24 May 2024 23:35:48 -0400
Subject: [PATCH 4/5] change for 2d arrays, use community sizes instead of
 memberships

---
 Cargo.lock                                    |   1 +
 rustworkx-core/Cargo.toml                     |   1 +
 rustworkx-core/src/generators/random_graph.rs | 165 ++++++++----------
 rustworkx/rustworkx.pyi                       |   8 +-
 src/random_graph.rs                           |  60 ++++---
 tests/test_random.py                          |  65 ++++---
 6 files changed, 151 insertions(+), 149 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 959c32034..25806f054 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -633,6 +633,7 @@ dependencies = [
  "fixedbitset",
  "hashbrown 0.14.5",
  "indexmap 2.2.6",
+ "ndarray",
  "num-traits",
  "petgraph",
  "priority-queue",
diff --git a/rustworkx-core/Cargo.toml b/rustworkx-core/Cargo.toml
index 781a9fbf5..c05a7f059 100644
--- a/rustworkx-core/Cargo.toml
+++ b/rustworkx-core/Cargo.toml
@@ -16,6 +16,7 @@ ahash.workspace = true
 fixedbitset.workspace = true
 hashbrown.workspace = true
 indexmap.workspace = true
+ndarray = "0.15.6"
 num-traits.workspace = true
 petgraph.workspace = true
 priority-queue = "2.0"
diff --git a/rustworkx-core/src/generators/random_graph.rs b/rustworkx-core/src/generators/random_graph.rs
index 189b058af..edea398fb 100644
--- a/rustworkx-core/src/generators/random_graph.rs
+++ b/rustworkx-core/src/generators/random_graph.rs
@@ -14,6 +14,7 @@
 
 use std::hash::Hash;
 
+use ndarray::ArrayView2;
 use petgraph::data::{Build, Create};
 use petgraph::visit::{
     Data, EdgeRef, GraphBase, GraphProp, IntoEdgeReferences, IntoEdgesDirected,
@@ -307,16 +308,17 @@ where
 
 /// Generate a graph from the stochastic block model.
 ///
-/// The stochastic block model is a generalization of the G<sub>np</sub> graph model
-/// (see [rustworkx_core::generators::gnp_random_graph] ). The connection probability of
-/// nodes `u` and `v` depends on their block (or community) and is given by
-/// `probabilities[blocks[u]][blocks[v]]`. The number of nodes and the number of blocks
-/// are inferred from `blocks`.
+/// The stochastic block model is a generalization of the G<sub>np</sub> random graph
+/// (see [gnp_random_graph] ). The connection probability of
+/// nodes `u` and `v` depends on their block and is given by
+/// `probabilities[blocks[u]][blocks[v]]`, where `blocks[u]` is the block membership
+/// of vertex `u`. The number of nodes and the number of blocks are inferred from
+/// `sizes`.
 ///
 /// Arguments:
 ///
-/// * `blocks` - Block membership (between 0 and B-1) of each node.
-/// * `probabilities` - B x B matrix that contains the connection probability between
+/// * `sizes` - Number of nodes in each block.
+/// * `probabilities` - B x B array that contains the connection probability between
 ///     nodes of different blocks. Must be symmetric for undirected graphs.
 /// * `loops` - Determines whether the graph can have loops or not.
 /// * `seed` - An optional seed to use for the random number generator.
@@ -327,12 +329,13 @@ where
 ///
 /// # Example
 /// ```rust
+/// use ndarray::arr2;
 /// use rustworkx_core::petgraph;
 /// use rustworkx_core::generators::sbm_random_graph;
 ///
 /// let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
-///     &vec![1, 0, 1],
-///     &[vec![0., 1.], vec![0., 1.]],
+///     &vec![1, 2],
+///     &ndarray::arr2(&[[0., 1.], [0., 1.]]).view(),
 ///     true,
 ///     Some(10),
 ///     || (),
@@ -343,8 +346,8 @@ where
 /// assert_eq!(g.edge_count(), 6);
 /// ```
 pub fn sbm_random_graph<G, T, F, H, M>(
-    blocks: &[usize],
-    probabilities: &[Vec<f64>],
+    sizes: &[usize],
+    probabilities: &ndarray::ArrayView2<f64>,
     loops: bool,
     seed: Option<u64>,
     mut default_node_weight: F,
@@ -356,27 +359,21 @@ where
     H: FnMut() -> M,
     G::NodeId: Eq + Hash,
 {
-    let num_nodes = blocks.len();
+    let num_nodes: usize = sizes.iter().sum();
     if num_nodes == 0 {
         return Err(InvalidInputError {});
     }
-    let num_communities = probabilities.len();
-    if probabilities
-        .iter()
-        .any(|xs| xs.len() != num_communities || xs.iter().any(|&x| !(0. ..=1.).contains(&x)))
+    let num_communities = sizes.len();
+    if probabilities.nrows() != num_communities
+        || probabilities.ncols() != num_communities
+        || probabilities.iter().any(|&x| !(0. ..=1.).contains(&x))
     {
         return Err(InvalidInputError {});
     }
-    if blocks.iter().max().unwrap_or(&usize::MAX) >= &num_communities {
-        return Err(InvalidInputError {});
-    }
-    if blocks.len() != num_nodes {
-        return Err(InvalidInputError {});
-    }
 
     let mut graph = G::with_capacity(num_nodes, num_nodes);
     let directed = graph.is_directed();
-    if !directed && !symmetric_matrix(probabilities) {
+    if !directed && !symmetric_array(probabilities) {
         return Err(InvalidInputError {});
     }
 
@@ -387,25 +384,30 @@ where
         Some(seed) => Pcg64::seed_from_u64(seed),
         None => Pcg64::from_entropy(),
     };
-    let between = Uniform::new(0.0, 1.0);
-
-    let mut block_partition: Vec<Vec<usize>> = (0..num_communities).map(|_| Vec::new()).collect();
-    for (node, block) in blocks.iter().enumerate() {
-        block_partition[*block].push(node);
+    let mut blocks = Vec::new();
+    {
+        let mut block = 0;
+        let mut vertices_left = sizes[0];
+        for _ in 0..num_nodes {
+            while vertices_left == 0 {
+                block += 1;
+                vertices_left = sizes[block];
+            }
+            blocks.push(block);
+            vertices_left -= 1;
+        }
     }
 
-    for (v, &b_v) in blocks.iter().enumerate().take(if directed || loops {
+    let between = Uniform::new(0.0, 1.0);
+    for v in 0..(if directed || loops {
         num_nodes
     } else {
         num_nodes - 1
     }) {
-        for (w, &b_w) in blocks
-            .iter()
-            .enumerate()
-            .skip(if directed { 0 } else { v })
-            .filter(|&(w, _)| w != v || loops)
-        {
-            if between.sample(&mut rng) < probabilities[b_v][b_w] {
+        for w in ((if directed { 0 } else { v })..num_nodes).filter(|&w| w != v || loops) {
+            if &between.sample(&mut rng)
+                < probabilities.get((blocks[v], blocks[w])).unwrap_or(&0_f64)
+            {
                 graph.add_edge(
                     graph.from_index(v),
                     graph.from_index(w),
@@ -417,14 +419,11 @@ where
     Ok(graph)
 }
 
-fn symmetric_matrix<T: std::cmp::PartialEq>(mat: &[Vec<T>]) -> bool {
-    let n = mat.len();
-    for (i, row) in mat.iter().enumerate().take(n - 1) {
-        if row.len() != n {
-            return false;
-        }
+fn symmetric_array<T: std::cmp::PartialEq>(mat: &ArrayView2<T>) -> bool {
+    let n = mat.nrows();
+    for (i, row) in mat.rows().into_iter().enumerate().take(n - 1) {
         for (j, m_ij) in row.iter().enumerate().skip(i + 1) {
-            if m_ij != &mat[j][i] {
+            if m_ij != mat.get((j, i)).unwrap() {
                 return false;
             }
         }
@@ -1007,31 +1006,11 @@ mod tests {
     }
 
     // Test sbm_random_graph
-    #[test]
-    fn test_sbm_directed_complete_blocks() {
-        let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![0., 1.]],
-            true,
-            Some(10),
-            || (),
-            || (),
-        )
-        .unwrap();
-        assert_eq!(g.node_count(), 3);
-        assert_eq!(g.edge_count(), 6);
-        for (u, v) in [(0, 0), (0, 2), (2, 0), (2, 2), (1, 0), (1, 2)] {
-            assert_eq!(g.contains_edge(u.into(), v.into()), true);
-        }
-        assert_eq!(g.contains_edge(0.into(), 1.into()), false);
-        assert_eq!(g.contains_edge(2.into(), 1.into()), false);
-    }
-
     #[test]
     fn test_sbm_directed_complete_blocks_loops() {
         let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![0., 1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1.], [0., 1.]]).view(),
             true,
             Some(10),
             || (),
@@ -1040,18 +1019,18 @@ mod tests {
         .unwrap();
         assert_eq!(g.node_count(), 3);
         assert_eq!(g.edge_count(), 6);
-        for (u, v) in [(0, 0), (0, 2), (2, 0), (2, 2), (1, 0), (1, 2)] {
+        for (u, v) in [(1, 1), (1, 2), (2, 1), (2, 2), (0, 1), (0, 2)] {
             assert_eq!(g.contains_edge(u.into(), v.into()), true);
         }
-        assert_eq!(g.contains_edge(0.into(), 1.into()), false);
-        assert_eq!(g.contains_edge(2.into(), 1.into()), false);
+        assert_eq!(g.contains_edge(1.into(), 0.into()), false);
+        assert_eq!(g.contains_edge(2.into(), 0.into()), false);
     }
 
     #[test]
     fn test_sbm_undirected_complete_blocks_loops() {
         let g = sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![1., 1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1.], [1., 1.]]).view(),
             true,
             Some(10),
             || (),
@@ -1060,17 +1039,17 @@ mod tests {
         .unwrap();
         assert_eq!(g.node_count(), 3);
         assert_eq!(g.edge_count(), 5);
-        for (u, v) in [(0, 0), (0, 2), (2, 2), (1, 0), (1, 2)] {
+        for (u, v) in [(1, 1), (1, 2), (2, 2), (0, 1), (0, 2)] {
             assert_eq!(g.contains_edge(u.into(), v.into()), true);
         }
-        assert_eq!(g.contains_edge(1.into(), 1.into()), false);
+        assert_eq!(g.contains_edge(0.into(), 0.into()), false);
     }
 
     #[test]
     fn test_sbm_directed_complete_blocks_noloops() {
         let g = sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![0., 1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1.], [0., 1.]]).view(),
             false,
             Some(10),
             || (),
@@ -1079,11 +1058,11 @@ mod tests {
         .unwrap();
         assert_eq!(g.node_count(), 3);
         assert_eq!(g.edge_count(), 4);
-        for (u, v) in [(0, 2), (2, 0), (1, 0), (1, 2)] {
+        for (u, v) in [(1, 2), (2, 1), (0, 1), (0, 2)] {
             assert_eq!(g.contains_edge(u.into(), v.into()), true);
         }
-        assert_eq!(g.contains_edge(0.into(), 1.into()), false);
-        assert_eq!(g.contains_edge(2.into(), 1.into()), false);
+        assert_eq!(g.contains_edge(1.into(), 0.into()), false);
+        assert_eq!(g.contains_edge(2.into(), 0.into()), false);
         for u in 0..2 {
             assert_eq!(g.contains_edge(u.into(), u.into()), false);
         }
@@ -1092,8 +1071,8 @@ mod tests {
     #[test]
     fn test_sbm_undirected_complete_blocks_noloops() {
         let g = sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![1., 1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1.], [1., 1.]]).view(),
             false,
             Some(10),
             || (),
@@ -1102,7 +1081,7 @@ mod tests {
         .unwrap();
         assert_eq!(g.node_count(), 3);
         assert_eq!(g.edge_count(), 3);
-        for (u, v) in [(0, 2), (1, 0), (1, 2)] {
+        for (u, v) in [(1, 2), (0, 1), (0, 2)] {
             assert_eq!(g.contains_edge(u.into(), v.into()), true);
         }
         for u in 0..2 {
@@ -1111,10 +1090,10 @@ mod tests {
     }
 
     #[test]
-    fn test_sbm_block_outofrange_error() {
+    fn test_sbm_bad_array_rows_error() {
         match sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 2],
-            &[vec![0., 1.], vec![1., 1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1.], [1., 1.], [1., 1.]]).view(),
             true,
             Some(10),
             || (),
@@ -1124,12 +1103,12 @@ mod tests {
             Err(e) => assert_eq!(e, InvalidInputError),
         };
     }
-
     #[test]
-    fn test_sbm_invalid_matrix_error() {
+
+    fn test_sbm_bad_array_cols_error() {
         match sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1., 1.], [1., 1., 1.]]).view(),
             true,
             Some(10),
             || (),
@@ -1141,10 +1120,10 @@ mod tests {
     }
 
     #[test]
-    fn test_sbm_asymmetric_matrix_error() {
+    fn test_sbm_asymmetric_array_error() {
         match sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![0., 1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1.], [0., 1.]]).view(),
             true,
             Some(10),
             || (),
@@ -1158,8 +1137,8 @@ mod tests {
     #[test]
     fn test_sbm_invalid_probability_error() {
         match sbm_random_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
-            &vec![1, 0, 1],
-            &[vec![0., 1.], vec![0., -1.]],
+            &vec![1, 2],
+            &ndarray::arr2(&[[0., 1.], [0., -1.]]).view(),
             true,
             Some(10),
             || (),
@@ -1174,7 +1153,7 @@ mod tests {
     fn test_sbm_empty_error() {
         match sbm_random_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
             &vec![],
-            &[],
+            &ndarray::arr2(&[[]]).view(),
             true,
             Some(10),
             || (),
diff --git a/rustworkx/rustworkx.pyi b/rustworkx/rustworkx.pyi
index c48b3f676..522962994 100644
--- a/rustworkx/rustworkx.pyi
+++ b/rustworkx/rustworkx.pyi
@@ -550,15 +550,15 @@ def undirected_gnp_random_graph(
     seed: int | None = ...,
 ) -> PyGraph: ...
 def directed_sbm_random_graph(
-    blocks: list[int],
-    probabilities: list[list[float]],
+    sizes: list[int],
+    probabilities: np.ndarray,
     loops: bool,
     /,
     seed: int | None = ...,
 ) -> PyDiGraph: ...
 def undirected_sbm_random_graph(
-    blocks: list[int],
-    probabilities: list[list[float]],
+    sizes: list[int],
+    probabilities: np.ndarray,
     loops: bool,
     /,
     seed: int | None = ...,
diff --git a/src/random_graph.rs b/src/random_graph.rs
index ce8857fd0..8360c0ff9 100644
--- a/src/random_graph.rs
+++ b/src/random_graph.rs
@@ -23,6 +23,8 @@ use petgraph::algo;
 use petgraph::graph::NodeIndex;
 use petgraph::prelude::*;
 
+use numpy::PyReadonlyArray2;
+
 use rand::distributions::{Distribution, Uniform};
 use rand::prelude::*;
 use rand_pcg::Pcg64;
@@ -275,37 +277,38 @@ pub fn undirected_gnm_random_graph(
 
 /// Return a directed graph from the stochastic block model.
 ///
-/// The stochastic block model is a generalization of the :math:`G(n,p)` graph model
-/// (see [rustworkx.undirected_gnp_random_graph] ). The connection probability of
+/// The stochastic block model is a generalization of the :math:`G(n,p)` random graph
+/// (see :func:`~rustworkx.directed_gnp_random_graph`). The connection probability of
 /// nodes ``u`` and ``v`` depends on their block (or community) and is given by
-/// ``probabilities[blocks[u]][blocks[v]]``. The number of nodes and the number of
-/// blocks are inferred from ``blocks``.
+/// ``probabilities[blocks[u]][blocks[v]]``, where ``blocks[u]`` is the block
+/// membership of node ``u``. The number of nodes and the number of blocks are
+/// inferred from ``sizes``.
 ///
 /// This algorithm has a time complexity of :math:`O(n^2)` for :math:`n` nodes.
 ///
 /// Arguments:
 ///
-/// :param list[int] blocks: Block membership (between 0 and B-1) of each node.
-/// :param list[list[float]] probabilities: B x B matrix that contains the
-///     connection probability between nodes of different blocks.
+/// :param list[int] sizes: Number of nodes in each block.
+/// :param np.ndarray probabilities: B x B array that contains the connection
+///     probability between nodes of different blocks.
 /// :param bool loops: Determines whether the graph can have loops or not.
 /// :param int seed:  An optional seed to use for the random number generator.
 ///
 /// :return: A PyDiGraph object
 /// :rtype: PyDiGraph
 #[pyfunction]
-#[pyo3(text_signature = "(blocks, probabilities, loops, /, seed=None)")]
-pub fn directed_sbm_random_graph(
-    py: Python,
-    blocks: Vec<usize>,
-    probabilities: Vec<Vec<f64>>,
+#[pyo3(text_signature = "(sizes, probabilities, loops, /, seed=None)")]
+pub fn directed_sbm_random_graph<'p>(
+    py: Python<'p>,
+    sizes: Vec<usize>,
+    probabilities: PyReadonlyArray2<'p, f64>,
     loops: bool,
     seed: Option<u64>,
 ) -> PyResult<digraph::PyDiGraph> {
     let default_fn = || py.None();
     let graph: StablePyGraph<Directed> = match core_generators::sbm_random_graph(
-        &blocks,
-        &probabilities,
+        &sizes,
+        &probabilities.as_array(),
         loops,
         seed,
         default_fn,
@@ -330,37 +333,38 @@ pub fn directed_sbm_random_graph(
 
 /// Return an undirected graph from the stochastic block model.
 ///
-/// The stochastic block model is a generalization of the :math:`G(n,p)` graph model
-/// (see [rustworkx.undirected_gnp_random_graph] ). The connection probability of
+/// The stochastic block model is a generalization of the :math:`G(n,p)` random graph
+/// (see :func:`~rustworkx.undirected_gnp_random_graph`). The connection probability of
 /// nodes ``u`` and ``v`` depends on their block (or community) and is given by
-/// ``probabilities[blocks[u]][blocks[v]]``. The number of nodes and the number of
-/// blocks are inferred from ``blocks``.
+/// ``probabilities[blocks[u]][blocks[v]]``, where ``blocks[u]`` is the block membership
+/// of node ``u``. The number of nodes and the number of blocks are inferred from
+/// ``sizes``.
 ///
 /// This algorithm has a time complexity of :math:`O(n^2)` for :math:`n` nodes.
 ///
 /// Arguments:
 ///
-/// :param list[int] blocks: Block membership (between 0 and B-1) of each node.
-/// :param list[list[float]] probabilities: Symmetric B x B matrix that contains
-///     the connection probability between nodes of different blocks.
+/// :param list[int] sizes: Number of nodes in each block.
+/// :param np.ndarray probabilities: Symmetric B x B array that contains the
+///     connection probability between nodes of different blocks.
 /// :param bool loops: Determines whether the graph can have loops or not.
 /// :param int seed:  An optional seed to use for the random number generator.
 ///
 /// :return: A PyGraph object
 /// :rtype: PyGraph
 #[pyfunction]
-#[pyo3(text_signature = "(blocks, probabilities, loops, /, seed=None)")]
-pub fn undirected_sbm_random_graph(
-    py: Python,
-    blocks: Vec<usize>,
-    probabilities: Vec<Vec<f64>>,
+#[pyo3(text_signature = "(sizes, probabilities, loops, /, seed=None)")]
+pub fn undirected_sbm_random_graph<'p>(
+    py: Python<'p>,
+    sizes: Vec<usize>,
+    probabilities: PyReadonlyArray2<'p, f64>,
     loops: bool,
     seed: Option<u64>,
 ) -> PyResult<graph::PyGraph> {
     let default_fn = || py.None();
     let graph: StablePyGraph<Undirected> = match core_generators::sbm_random_graph(
-        &blocks,
-        &probabilities,
+        &sizes,
+        &probabilities.as_array(),
         loops,
         seed,
         default_fn,
diff --git a/tests/test_random.py b/tests/test_random.py
index f7d1267b5..74f7668bb 100644
--- a/tests/test_random.py
+++ b/tests/test_random.py
@@ -14,6 +14,7 @@
 import random
 import math
 
+import numpy as np
 import rustworkx
 
 
@@ -179,23 +180,29 @@ def test_random_gnm_undirected_payload(self):
 
 class TestRandomSBM(unittest.TestCase):
     def test_undirected_sbm_complete_blocks_loops(self):
-        graph = rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 1], [1, 0]], True)
+        graph = rustworkx.undirected_sbm_random_graph(
+            [2, 1], np.array([[1, 1], [1, 0]], dtype=float), True
+        )
         self.assertEqual(len(graph), 3)
         self.assertEqual(len(graph.edges()), 5)
         for i in range(2):
             for j in range(i, 2):
-                if (i, j) != (1, 1):
+                if (i, j) != (2, 2):
                     self.assertTrue(graph.has_edge(i, j))
-        self.assertFalse(graph.has_edge(1, 1))
+        self.assertFalse(graph.has_edge(2, 2))
 
     def test_directed_sbm_complete_blocks_loops(self):
-        graph = rustworkx.directed_sbm_random_graph([0, 1, 0], [[0, 0], [1, 1]], True)
+        graph = rustworkx.directed_sbm_random_graph(
+            [2, 1], np.array([[0, 0], [1, 1]], dtype=float), True
+        )
         self.assertEqual(len(graph), 3)
         self.assertEqual(len(graph.edges()), 3)
-        self.assertEqual(set(graph.edge_list()), set([(1, 1), (1, 0), (1, 2)]))
+        self.assertEqual(set(graph.edge_list()), set([(2, 2), (2, 0), (2, 1)]))
 
     def test_undirected_sbm_complete_blocks_noloops(self):
-        graph = rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 1], [1, 0]], False)
+        graph = rustworkx.undirected_sbm_random_graph(
+            [2, 1], np.array([[1, 1], [1, 0]], dtype=float), False
+        )
         self.assertEqual(len(graph), 3)
         self.assertEqual(len(graph.edges()), 3)
         for i in range(2):
@@ -204,42 +211,52 @@ def test_undirected_sbm_complete_blocks_noloops(self):
                     self.assertTrue(graph.has_edge(i, j))
 
     def test_directed_sbm_complete_blocks_noloops(self):
-        graph = rustworkx.directed_sbm_random_graph([0, 1, 0], [[0, 0], [1, 1]], False)
+        graph = rustworkx.directed_sbm_random_graph(
+            [2, 1], np.array([[0, 0], [1, 1]], dtype=float), False
+        )
         self.assertEqual(len(graph), 3)
         self.assertEqual(len(graph.edges()), 2)
-        self.assertEqual(set(graph.edge_list()), set([(1, 0), (1, 2)]))
+        self.assertEqual(set(graph.edge_list()), set([(2, 0), (2, 1)]))
 
     def test_undirected_sbm_asymmetric_probabilities_error(self):
         with self.assertRaises(ValueError):
-            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[0, 0], [1, 1]], True)
-
-    def test_sbm_out_of_range_blocks_error(self):
-        with self.assertRaises(ValueError):
-            rustworkx.undirected_sbm_random_graph([0, 2, 0], [[1, 0], [0, 1]], True)
-        with self.assertRaises(ValueError):
-            rustworkx.directed_sbm_random_graph([0, 2, 0], [[1, 0], [0, 1]], True)
+            rustworkx.undirected_sbm_random_graph(
+                [2, 1], np.array([[0, 0], [1, 1]], dtype=float), True
+            )
 
-    def test_sbm_invalid_matrix(self):
+    def test_sbm_invalid_matrix_dim(self):
         with self.assertRaises(ValueError):
-            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1, 0]], True)
+            rustworkx.undirected_sbm_random_graph(
+                [2, 1], np.array([[1, 0], [0, 1], [0, 1]], dtype=float), True
+            )
         with self.assertRaises(ValueError):
-            rustworkx.directed_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1, 0]], True)
+            rustworkx.directed_sbm_random_graph(
+                [2, 1], np.array([[1, 0, 1], [0, 1, 0]], dtype=float), True
+            )
 
     def test_sbm_invalid_probabilities(self):
         with self.assertRaises(ValueError):
-            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1.5]], True)
+            rustworkx.undirected_sbm_random_graph(
+                [2, 1], np.array([[1, 0], [0, 1.5]], dtype=float), True
+            )
         with self.assertRaises(ValueError):
-            rustworkx.undirected_sbm_random_graph([0, 1, 0], [[-1, 0], [0, 1]], True)
+            rustworkx.undirected_sbm_random_graph(
+                [2, 1], np.array([[-1, 0], [0, 1]], dtype=float), True
+            )
         with self.assertRaises(ValueError):
-            rustworkx.directed_sbm_random_graph([0, 1, 0], [[1, 0], [0, 1.5]], True)
+            rustworkx.directed_sbm_random_graph(
+                [2, 1], np.array([[1, 0], [0, 1.5]], dtype=float), True
+            )
         with self.assertRaises(ValueError):
-            rustworkx.directed_sbm_random_graph([0, 1, 0], [[-1, 0], [0, 1]], True)
+            rustworkx.directed_sbm_random_graph(
+                [2, 1], np.array([[-1, 0], [0, 1]], dtype=float), True
+            )
 
     def test_sbm_empty(self):
         with self.assertRaises(ValueError):
-            rustworkx.undirected_sbm_random_graph([], [], True)
+            rustworkx.undirected_sbm_random_graph([], np.array([[]]), True)
         with self.assertRaises(ValueError):
-            rustworkx.directed_sbm_random_graph([], [], True)
+            rustworkx.directed_sbm_random_graph([], np.array([[]]), True)
 
 
 class TestGeometricRandomGraph(unittest.TestCase):

From 8d86da1639b606754fb4af877faa43830a7ae3cd Mon Sep 17 00:00:00 2001
From: Ivan Carvalho <ivancarv@student.ubc.ca>
Date: Sat, 25 May 2024 17:29:40 -0400
Subject: [PATCH 5/5] Use cargo workspace for ndarray deps

---
 Cargo.toml                | 6 ++----
 rustworkx-core/Cargo.toml | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 1e50fea81..c3bb08693 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,6 +28,7 @@ ahash = "0.8.6"
 fixedbitset = "0.4.2"
 hashbrown = { version = ">=0.13, <0.15", features = ["rayon"] }
 indexmap = { version = ">=1.9, <3", features = ["rayon"] }
+ndarray = { version = "0.15.6", features = ["rayon"] }
 num-traits = "0.2"
 numpy = "0.21.0"
 petgraph = "0.6.5"
@@ -44,6 +45,7 @@ ahash.workspace = true
 fixedbitset.workspace = true
 hashbrown.workspace = true
 indexmap.workspace = true
+ndarray.workspace = true
 ndarray-stats = "0.5.1"
 num-bigint = "0.4"
 num-complex = "0.4"
@@ -63,10 +65,6 @@ rustworkx-core = { path = "rustworkx-core", version = "=0.15.0" }
 version = "0.21.2"
 features = ["abi3-py38", "extension-module", "hashbrown", "num-bigint", "num-complex", "indexmap"]
 
-[dependencies.ndarray]
-version = "^0.15.6"
-features = ["rayon"]
-
 [dependencies.sprs]
 version = "^0.11"
 features = ["multi_thread"]
diff --git a/rustworkx-core/Cargo.toml b/rustworkx-core/Cargo.toml
index c05a7f059..c8d292627 100644
--- a/rustworkx-core/Cargo.toml
+++ b/rustworkx-core/Cargo.toml
@@ -16,7 +16,7 @@ ahash.workspace = true
 fixedbitset.workspace = true
 hashbrown.workspace = true
 indexmap.workspace = true
-ndarray = "0.15.6"
+ndarray.workspace = true
 num-traits.workspace = true
 petgraph.workspace = true
 priority-queue = "2.0"