diff --git a/docs/source/api/generators.rst b/docs/source/api/generators.rst index 2f930704c7..ad70e6c7fd 100644 --- a/docs/source/api/generators.rst +++ b/docs/source/api/generators.rst @@ -16,6 +16,8 @@ Generators rustworkx.generators.directed_mesh_graph rustworkx.generators.grid_graph rustworkx.generators.directed_grid_graph + rustworkx.generators.nd_grid_graph + rustworkx.generators.directed_nd_grid_graph rustworkx.generators.binomial_tree_graph rustworkx.generators.directed_binomial_tree_graph rustworkx.generators.hexagonal_lattice_graph diff --git a/releasenotes/notes/nd-grid-graph-7f3a9c2e1b8d.yaml b/releasenotes/notes/nd-grid-graph-7f3a9c2e1b8d.yaml new file mode 100644 index 0000000000..f30eba5e7b --- /dev/null +++ b/releasenotes/notes/nd-grid-graph-7f3a9c2e1b8d.yaml @@ -0,0 +1,19 @@ +--- +features: + - | + Added :func:`~rustworkx.generators.nd_grid_graph` and + :func:`~rustworkx.generators.directed_nd_grid_graph` for generating + n-dimensional grid graphs. These extend the existing 2D grid_graph + to support arbitrary dimensions. + + .. jupyter-execute:: + + import rustworkx.generators + + # 3D cube + cube = rustworkx.generators.nd_grid_graph([2, 2, 2]) + print(f"{len(cube)} nodes, {len(cube.edges())} edges") + + # with coordinate labels + grid = rustworkx.generators.nd_grid_graph([3, 3], with_positions=True) + print(f"node 0: {grid[0]}") diff --git a/rustworkx-core/src/generators/mod.rs b/rustworkx-core/src/generators/mod.rs index 8133b157ac..867b79b99f 100644 --- a/rustworkx-core/src/generators/mod.rs +++ b/rustworkx-core/src/generators/mod.rs @@ -22,6 +22,7 @@ mod grid_graph; mod heavy_hex_graph; mod heavy_square_graph; mod hexagonal_lattice_graph; +mod nd_grid_graph; mod karate_club; mod lollipop_graph; mod path_graph; @@ -58,6 +59,7 @@ pub use heavy_square_graph::heavy_square_graph; pub use hexagonal_lattice_graph::{hexagonal_lattice_graph, hexagonal_lattice_graph_weighted}; pub use karate_club::karate_club_graph; pub use lollipop_graph::lollipop_graph; +pub use nd_grid_graph::{nd_grid_graph, nd_grid_graph_weighted}; pub use path_graph::path_graph; pub use petersen_graph::petersen_graph; pub use random_graph::barabasi_albert_graph; diff --git a/rustworkx-core/src/generators/nd_grid_graph.rs b/rustworkx-core/src/generators/nd_grid_graph.rs new file mode 100644 index 0000000000..d1507ad88f --- /dev/null +++ b/rustworkx-core/src/generators/nd_grid_graph.rs @@ -0,0 +1,320 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); you may +// not use this file except in compliance with the License. You may obtain +// a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +// License for the specific language governing permissions and limitations +// under the License. + +use petgraph::data::{Build, Create}; +use petgraph::visit::{Data, NodeIndexable}; + +use super::InvalidInputError; + +// Helper to convert flat index -> coordinates +fn index_to_coords(index: usize, dim: &[usize]) -> Vec { + let mut coords = Vec::with_capacity(dim.len()); + let mut rem = index; + for &d in dim { + coords.push(rem % d); + rem /= d; + } + coords +} + +// Helper to convert coordinates -> flat index +fn coords_to_index(coords: &[usize], dim: &[usize]) -> usize { + let mut idx = 0; + let mut mult = 1; + for (i, &c) in coords.iter().enumerate() { + idx += c * mult; + mult *= dim[i]; + } + idx +} + +/// Generate an n-dimensional grid graph. +/// +/// Nodes are placed at integer coordinates in an n-dimensional box, and edges +/// connect nodes that differ by 1 in exactly one coordinate. +/// +/// # Example +/// ```rust +/// use rustworkx_core::petgraph; +/// use rustworkx_core::generators::nd_grid_graph; +/// +/// let g: petgraph::graph::UnGraph<(), ()> = nd_grid_graph( +/// &[2, 2, 2], || (), || (), false, None +/// ).unwrap(); +/// assert_eq!(g.node_count(), 8); +/// assert_eq!(g.edge_count(), 12); +/// ``` +pub fn nd_grid_graph( + dim: &[usize], + mut default_node_weight: F, + mut default_edge_weight: H, + bidirectional: bool, + periodic: Option<&[bool]>, +) -> Result +where + G: Build + Create + Data + NodeIndexable, + F: FnMut() -> T, + H: FnMut() -> M, +{ + if dim.is_empty() || dim.iter().any(|&d| d == 0) { + return Err(InvalidInputError {}); + } + + if let Some(p) = periodic { + if p.len() != dim.len() { + return Err(InvalidInputError {}); + } + // periodic dim needs at least 2 nodes (can't wrap to self) + for (i, &is_per) in p.iter().enumerate() { + if is_per && dim[i] < 2 { + return Err(InvalidInputError {}); + } + } + } + + let num_nodes: usize = dim.iter().product(); + + // count edges: each dimension contributes (size-1) edges per line, or size if periodic + let mut num_edges = 0; + for (i, &sz) in dim.iter().enumerate() { + let per_line = if periodic.map_or(false, |p| p[i]) { sz } else { sz - 1 }; + num_edges += per_line * (num_nodes / sz); + } + if bidirectional { + num_edges *= 2; + } + + let mut graph = G::with_capacity(num_nodes, num_edges); + + for _ in 0..num_nodes { + graph.add_node(default_node_weight()); + } + + // add edges: connect each node to its neighbor in positive direction per dimension + for node_idx in 0..num_nodes { + let coords = index_to_coords(node_idx, dim); + + for (d, &sz) in dim.iter().enumerate() { + let is_per = periodic.map_or(false, |p| p[d]); + + if coords[d] + 1 < sz { + let mut nb = coords.clone(); + nb[d] += 1; + let nb_idx = coords_to_index(&nb, dim); + + let a = graph.from_index(node_idx); + let b = graph.from_index(nb_idx); + graph.add_edge(a, b, default_edge_weight()); + if bidirectional { + graph.add_edge(b, a, default_edge_weight()); + } + } else if is_per { + // wrap around + let mut nb = coords.clone(); + nb[d] = 0; + let nb_idx = coords_to_index(&nb, dim); + + let a = graph.from_index(node_idx); + let b = graph.from_index(nb_idx); + graph.add_edge(a, b, default_edge_weight()); + if bidirectional { + graph.add_edge(b, a, default_edge_weight()); + } + } + } + } + + Ok(graph) +} + +/// Like `nd_grid_graph` but assigns node weights based on coordinates. +pub fn nd_grid_graph_weighted( + dim: &[usize], + mut node_weight_fn: F, + mut default_edge_weight: H, + bidirectional: bool, + periodic: Option<&[bool]>, +) -> Result +where + G: Build + Create + Data + NodeIndexable, + F: FnMut(&[usize]) -> T, + H: FnMut() -> M, +{ + if dim.is_empty() || dim.iter().any(|&d| d == 0) { + return Err(InvalidInputError {}); + } + + if let Some(p) = periodic { + if p.len() != dim.len() { + return Err(InvalidInputError {}); + } + for (i, &is_per) in p.iter().enumerate() { + if is_per && dim[i] < 2 { + return Err(InvalidInputError {}); + } + } + } + + let num_nodes: usize = dim.iter().product(); + + let mut num_edges = 0; + for (i, &sz) in dim.iter().enumerate() { + let per_line = if periodic.map_or(false, |p| p[i]) { sz } else { sz - 1 }; + num_edges += per_line * (num_nodes / sz); + } + if bidirectional { + num_edges *= 2; + } + + let mut graph = G::with_capacity(num_nodes, num_edges); + + for i in 0..num_nodes { + let coords = index_to_coords(i, dim); + graph.add_node(node_weight_fn(&coords)); + } + + for node_idx in 0..num_nodes { + let coords = index_to_coords(node_idx, dim); + + for (d, &sz) in dim.iter().enumerate() { + let is_per = periodic.map_or(false, |p| p[d]); + + if coords[d] + 1 < sz { + let mut nb = coords.clone(); + nb[d] += 1; + let nb_idx = coords_to_index(&nb, dim); + + let a = graph.from_index(node_idx); + let b = graph.from_index(nb_idx); + graph.add_edge(a, b, default_edge_weight()); + if bidirectional { + graph.add_edge(b, a, default_edge_weight()); + } + } else if is_per { + let mut nb = coords.clone(); + nb[d] = 0; + let nb_idx = coords_to_index(&nb, dim); + + let a = graph.from_index(node_idx); + let b = graph.from_index(nb_idx); + graph.add_edge(a, b, default_edge_weight()); + if bidirectional { + graph.add_edge(b, a, default_edge_weight()); + } + } + } + } + + Ok(graph) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::petgraph; + use petgraph::visit::EdgeRef; + + #[test] + fn test_line_graph() { + let g: petgraph::graph::UnGraph<(), ()> = + nd_grid_graph(&[5], || (), || (), false, None).unwrap(); + assert_eq!(g.node_count(), 5); + assert_eq!(g.edge_count(), 4); + } + + #[test] + fn test_grid_2d() { + let g: petgraph::graph::UnGraph<(), ()> = + nd_grid_graph(&[3, 4], || (), || (), false, None).unwrap(); + assert_eq!(g.node_count(), 12); + assert_eq!(g.edge_count(), 17); + } + + #[test] + fn test_cube() { + let g: petgraph::graph::UnGraph<(), ()> = + nd_grid_graph(&[2, 2, 2], || (), || (), false, None).unwrap(); + assert_eq!(g.node_count(), 8); + assert_eq!(g.edge_count(), 12); + } + + #[test] + fn test_hypercube_4d() { + let g: petgraph::graph::UnGraph<(), ()> = + nd_grid_graph(&[2, 2, 2, 2], || (), || (), false, None).unwrap(); + assert_eq!(g.node_count(), 16); + assert_eq!(g.edge_count(), 32); + } + + #[test] + fn test_bidir() { + let g: petgraph::graph::DiGraph<(), ()> = + nd_grid_graph(&[2, 2], || (), || (), true, None).unwrap(); + assert_eq!(g.edge_count(), 8); + } + + #[test] + fn test_cycle() { + let g: petgraph::graph::UnGraph<(), ()> = + nd_grid_graph(&[5], || (), || (), false, Some(&[true])).unwrap(); + assert_eq!(g.edge_count(), 5); + } + + #[test] + fn test_torus() { + let g: petgraph::graph::UnGraph<(), ()> = + nd_grid_graph(&[3, 3], || (), || (), false, Some(&[true, true])).unwrap(); + assert_eq!(g.edge_count(), 18); + } + + #[test] + fn test_coords() { + let g: petgraph::graph::UnGraph, ()> = + nd_grid_graph_weighted(&[2, 3], |c| c.to_vec(), || (), false, None).unwrap(); + assert_eq!(g[petgraph::graph::NodeIndex::new(0)], vec![0, 0]); + assert_eq!(g[petgraph::graph::NodeIndex::new(1)], vec![1, 0]); + assert_eq!(g[petgraph::graph::NodeIndex::new(3)], vec![1, 1]); + } + + #[test] + fn test_empty_err() { + let r: Result, _> = + nd_grid_graph(&[], || (), || (), false, None); + assert!(r.is_err()); + } + + #[test] + fn test_zero_err() { + let r: Result, _> = + nd_grid_graph(&[2, 0], || (), || (), false, None); + assert!(r.is_err()); + } + + #[test] + fn test_cube_edges() { + let g: petgraph::graph::UnGraph<(), ()> = + nd_grid_graph(&[2, 2, 2], || (), || (), false, None).unwrap(); + + let edges: Vec<(usize, usize)> = g + .edge_references() + .map(|e| { + let (a, b) = (e.source().index(), e.target().index()); + (a.min(b), a.max(b)) + }) + .collect(); + + // check some expected edges + assert!(edges.contains(&(0, 1))); + assert!(edges.contains(&(0, 2))); + assert!(edges.contains(&(0, 4))); + } +} diff --git a/rustworkx/generators/__init__.pyi b/rustworkx/generators/__init__.pyi index 6136cb3de1..e08471fb4f 100644 --- a/rustworkx/generators/__init__.pyi +++ b/rustworkx/generators/__init__.pyi @@ -64,6 +64,19 @@ def directed_grid_graph( bidirectional: bool = ..., multigraph: bool = ..., ) -> PyDiGraph: ... +def nd_grid_graph( + dim: Sequence[int], + periodic: bool = ..., + with_positions: bool = ..., + multigraph: bool = ..., +) -> PyGraph: ... +def directed_nd_grid_graph( + dim: Sequence[int], + periodic: bool = ..., + with_positions: bool = ..., + bidirectional: bool = ..., + multigraph: bool = ..., +) -> PyDiGraph: ... def heavy_square_graph(d: int, multigraph: bool = ...) -> PyGraph: ... def directed_heavy_square_graph( d: int, diff --git a/src/generators.rs b/src/generators.rs index d8cb699bd9..5ad3afe3ee 100644 --- a/src/generators.rs +++ b/src/generators.rs @@ -592,6 +592,198 @@ pub fn directed_grid_graph( }) } +/// Generate an undirected n-dimensional grid graph. +/// +/// The n-dimensional grid graph has nodes at each point `(i_0, i_1, ..., i_{n-1})` +/// where `0 <= i_k < dim[k]` for each dimension k. Two nodes are connected by an +/// edge if they differ by exactly 1 in exactly one coordinate. +/// +/// :param list dim: A sequence of integers representing the size of each dimension. +/// For example, ``[2, 3, 4]`` creates a 3D grid of size 2×3×4 with 24 nodes. +/// :param bool periodic: When set to ``True``, all dimensions will be periodic +/// (wrapping), creating a torus-like structure. Defaults to ``False``. +/// :param bool with_positions: When set to ``True``, each node will be assigned +/// a tuple containing its coordinates in the grid as its weight. Defaults to ``False``. +/// :param bool multigraph: When set to ``False``, the output +/// :class:`~rustworkx.PyGraph` object will not be a multigraph and +/// won't allow parallel edges to be added. Instead, +/// calls which would create a parallel edge will update the existing edge. +/// +/// :returns: The generated n-dimensional grid graph +/// :rtype: PyGraph +/// :raises IndexError: If ``dim`` is empty or contains zeros +/// :raises ValueError: If periodic is enabled for a dimension with size < 2 +/// +/// .. jupyter-execute:: +/// +/// import rustworkx.generators +/// from rustworkx.visualization import mpl_draw +/// +/// # Create a 3x3x3 cube +/// graph = rustworkx.generators.nd_grid_graph([3, 3, 3]) +/// len(graph) # 27 nodes +/// +#[pyfunction] +#[pyo3( + signature=(dim, periodic=false, with_positions=false, multigraph=true), +)] +pub fn nd_grid_graph( + py: Python, + dim: Vec, + periodic: bool, + with_positions: bool, + multigraph: bool, +) -> PyResult { + let periodic_vec: Option> = if periodic { + Some(vec![true; dim.len()]) + } else { + None + }; + + let graph: StablePyGraph = if with_positions { + let node_weight_fn = |coords: &[usize]| { + let tuple: Vec = coords.to_vec(); + tuple.into_py_any(py).unwrap() + }; + match core_generators::nd_grid_graph_weighted( + &dim, + node_weight_fn, + || py.None(), + false, + periodic_vec.as_deref(), + ) { + Ok(graph) => graph, + Err(_) => { + return Err(PyIndexError::new_err( + "Invalid dimensions: dim must be non-empty and contain no zeros", + )) + } + } + } else { + match core_generators::nd_grid_graph( + &dim, + || py.None(), + || py.None(), + false, + periodic_vec.as_deref(), + ) { + Ok(graph) => graph, + Err(_) => { + return Err(PyIndexError::new_err( + "Invalid dimensions: dim must be non-empty and contain no zeros", + )) + } + } + }; + + Ok(graph::PyGraph { + graph, + node_removed: false, + multigraph, + attrs: py.None(), + }) +} + +/// Generate a directed n-dimensional grid graph. +/// +/// The n-dimensional grid graph has nodes at each point `(i_0, i_1, ..., i_{n-1})` +/// where `0 <= i_k < dim[k]` for each dimension k. Two nodes are connected by an +/// edge if they differ by exactly 1 in exactly one coordinate. +/// +/// The edges propagate in the positive direction along each axis if ``bidirectional`` +/// is ``False``. +/// +/// :param list dim: A sequence of integers representing the size of each dimension. +/// For example, ``[2, 3, 4]`` creates a 3D grid of size 2×3×4 with 24 nodes. +/// :param bool periodic: When set to ``True``, all dimensions will be periodic +/// (wrapping), creating a torus-like structure. Defaults to ``False``. +/// :param bool with_positions: When set to ``True``, each node will be assigned +/// a tuple containing its coordinates in the grid as its weight. Defaults to ``False``. +/// :param bool bidirectional: A parameter to indicate if edges should exist in +/// both directions between nodes. Defaults to ``False``. +/// :param bool multigraph: When set to ``False``, the output +/// :class:`~rustworkx.PyDiGraph` object will not be a multigraph and +/// won't allow parallel edges to be added. Instead, +/// calls which would create a parallel edge will update the existing edge. +/// +/// :returns: The generated directed n-dimensional grid graph +/// :rtype: PyDiGraph +/// :raises IndexError: If ``dim`` is empty or contains zeros +/// :raises ValueError: If periodic is enabled for a dimension with size < 2 +/// +/// .. jupyter-execute:: +/// +/// import rustworkx.generators +/// from rustworkx.visualization import mpl_draw +/// +/// # Create a directed 2x2x2 cube +/// graph = rustworkx.generators.directed_nd_grid_graph([2, 2, 2]) +/// len(graph) # 8 nodes +/// +#[pyfunction] +#[pyo3( + signature=(dim, periodic=false, with_positions=false, bidirectional=false, multigraph=true), +)] +pub fn directed_nd_grid_graph( + py: Python, + dim: Vec, + periodic: bool, + with_positions: bool, + bidirectional: bool, + multigraph: bool, +) -> PyResult { + let periodic_vec: Option> = if periodic { + Some(vec![true; dim.len()]) + } else { + None + }; + + let graph: StablePyGraph = if with_positions { + let node_weight_fn = |coords: &[usize]| { + let tuple: Vec = coords.to_vec(); + tuple.into_py_any(py).unwrap() + }; + match core_generators::nd_grid_graph_weighted( + &dim, + node_weight_fn, + || py.None(), + bidirectional, + periodic_vec.as_deref(), + ) { + Ok(graph) => graph, + Err(_) => { + return Err(PyIndexError::new_err( + "Invalid dimensions: dim must be non-empty and contain no zeros", + )) + } + } + } else { + match core_generators::nd_grid_graph( + &dim, + || py.None(), + || py.None(), + bidirectional, + periodic_vec.as_deref(), + ) { + Ok(graph) => graph, + Err(_) => { + return Err(PyIndexError::new_err( + "Invalid dimensions: dim must be non-empty and contain no zeros", + )) + } + } + }; + + Ok(digraph::PyDiGraph { + graph, + node_removed: false, + check_cycle: false, + cycle_state: algo::DfsSpace::default(), + multigraph, + attrs: py.None(), + }) +} + /// Generate an undirected heavy square graph. /// /// Fig. 6 of https://arxiv.org/abs/1907.09528. @@ -1790,6 +1982,8 @@ pub fn generators(_py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(directed_mesh_graph))?; m.add_wrapped(wrap_pyfunction!(grid_graph))?; m.add_wrapped(wrap_pyfunction!(directed_grid_graph))?; + m.add_wrapped(wrap_pyfunction!(nd_grid_graph))?; + m.add_wrapped(wrap_pyfunction!(directed_nd_grid_graph))?; m.add_wrapped(wrap_pyfunction!(heavy_square_graph))?; m.add_wrapped(wrap_pyfunction!(directed_heavy_square_graph))?; m.add_wrapped(wrap_pyfunction!(heavy_hex_graph))?; diff --git a/tests/generators/test_grid.py b/tests/generators/test_grid.py index 231ab57110..ebb8c3a174 100644 --- a/tests/generators/test_grid.py +++ b/tests/generators/test_grid.py @@ -129,3 +129,87 @@ def test_grid_no_weights_or_dim(self): rustworkx.generators.grid_graph() rustworkx.generators.grid_graph(rows=5, weights=[1] * 5) rustworkx.generators.grid_graph(cols=5, weights=[1] * 5) + + +class TestNdGridGraph(unittest.TestCase): + def test_1d_grid(self): + graph = rustworkx.generators.nd_grid_graph([5]) + self.assertEqual(len(graph), 5) + self.assertEqual(len(graph.edges()), 4) + + def test_2d_grid(self): + graph = rustworkx.generators.nd_grid_graph([3, 4]) + self.assertEqual(len(graph), 12) + self.assertEqual(len(graph.edges()), 17) + + def test_3d_cube(self): + graph = rustworkx.generators.nd_grid_graph([2, 2, 2]) + self.assertEqual(len(graph), 8) + self.assertEqual(len(graph.edges()), 12) + + def test_3d_larger(self): + graph = rustworkx.generators.nd_grid_graph([2, 3, 4]) + self.assertEqual(len(graph), 24) + self.assertEqual(len(graph.edges()), 46) + + def test_4d(self): + graph = rustworkx.generators.nd_grid_graph([2, 2, 2, 2]) + self.assertEqual(len(graph), 16) + self.assertEqual(len(graph.edges()), 32) + + def test_with_positions(self): + graph = rustworkx.generators.nd_grid_graph([2, 3], with_positions=True) + nodes = graph.nodes() + self.assertEqual(nodes[0], [0, 0]) + self.assertEqual(nodes[1], [1, 0]) + self.assertEqual(nodes[3], [1, 1]) + + def test_periodic(self): + # 1D periodic is a cycle + graph = rustworkx.generators.nd_grid_graph([5], periodic=True) + self.assertEqual(len(graph.edges()), 5) + + def test_torus(self): + graph = rustworkx.generators.nd_grid_graph([3, 3], periodic=True) + self.assertEqual(len(graph.edges()), 18) + + def test_empty_dim(self): + with self.assertRaises(IndexError): + rustworkx.generators.nd_grid_graph([]) + + def test_zero_dim(self): + with self.assertRaises(IndexError): + rustworkx.generators.nd_grid_graph([2, 0, 3]) + + +class TestDirectedNdGridGraph(unittest.TestCase): + def test_basic(self): + graph = rustworkx.generators.directed_nd_grid_graph([3, 3]) + self.assertEqual(len(graph), 9) + self.assertEqual(len(graph.edges()), 12) + + def test_cube(self): + graph = rustworkx.generators.directed_nd_grid_graph([2, 2, 2]) + self.assertEqual(len(graph), 8) + self.assertEqual(len(graph.edges()), 12) + + def test_bidirectional(self): + graph = rustworkx.generators.directed_nd_grid_graph([2, 2, 2], bidirectional=True) + self.assertEqual(len(graph.edges()), 24) + + def test_with_positions(self): + graph = rustworkx.generators.directed_nd_grid_graph([2, 2], with_positions=True) + self.assertEqual(graph.nodes()[0], [0, 0]) + self.assertEqual(graph.nodes()[3], [1, 1]) + + def test_periodic(self): + graph = rustworkx.generators.directed_nd_grid_graph([3, 3], periodic=True) + self.assertEqual(len(graph.edges()), 18) + + def test_edge_direction(self): + graph = rustworkx.generators.directed_nd_grid_graph([2, 2]) + out_edges = graph.out_edges(0) + out_targets = [e[1] for e in out_edges] + self.assertIn(1, out_targets) + self.assertIn(2, out_targets) + self.assertEqual(len(graph.in_edges(0)), 0)