diff --git a/docs/source/api.rst b/docs/source/api.rst index 5d4397688..bdc349b5a 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -30,6 +30,7 @@ Generators retworkx.generators.directed_mesh_graph retworkx.generators.grid_graph retworkx.generators.directed_grid_graph + retworkx.generators.binomial_tree_graph retworkx.generators.hexagonal_lattice_graph retworkx.generators.directed_hexagonal_lattice_graph diff --git a/releasenotes/notes/add-binomial-tree-graph-generator-1f6ff6ba3809b901.yaml b/releasenotes/notes/add-binomial-tree-graph-generator-1f6ff6ba3809b901.yaml new file mode 100644 index 000000000..23bef73d6 --- /dev/null +++ b/releasenotes/notes/add-binomial-tree-graph-generator-1f6ff6ba3809b901.yaml @@ -0,0 +1,4 @@ +--- +features: + - | + Added a new generator for constructing a binomial tree graph (:func:`retworkx.generators.binomial_tree_graph`). diff --git a/src/generators.rs b/src/generators.rs index f3fd483b2..8750d631f 100644 --- a/src/generators.rs +++ b/src/generators.rs @@ -15,6 +15,7 @@ use std::iter; use petgraph::algo; use petgraph::graph::NodeIndex; use petgraph::stable_graph::{StableDiGraph, StableUnGraph}; +use petgraph::visit::{EdgeRef, IntoEdgeReferences}; use pyo3::exceptions::PyIndexError; use pyo3::prelude::*; @@ -983,6 +984,247 @@ pub fn directed_grid_graph( }) } +/// Generate an undirected binomial tree of order n recursively. +/// +/// :param int order: Order of the binomial tree. +/// :param list weights: A list of node weights. If the number of weights is +/// less than 2**order extra nodes with with None will be appended. +/// :param bool multigraph: When set to False the output +/// :class:`~retworkx.PyGraph` object will not be 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: A binomial tree with 2^n vertices and 2^n - 1 edges. +/// :rtype: PyGraph +/// :raises IndexError: If the lenght of ``weights`` is greater that 2^n +/// +/// .. jupyter-execute:: +/// +/// import os +/// import tempfile +/// +/// import pydot +/// from PIL import Image +/// +/// import retworkx.generators +/// +/// graph = retworkx.generators.binomial_tree_graph(4) +/// dot_str = graph.to_dot( +/// lambda node: dict( +/// color='black', fillcolor='lightblue', style='filled')) +/// dot = pydot.graph_from_dot_data(dot_str)[0] +/// +/// with tempfile.TemporaryDirectory() as tmpdirname: +/// tmp_path = os.path.join(tmpdirname, 'dag.png') +/// dot.write_png(tmp_path) +/// image = Image.open(tmp_path) +/// os.remove(tmp_path) +/// image +/// +#[pyfunction(multigraph = true)] +#[text_signature = "(order, /, weights=None, multigraph=True)"] +pub fn binomial_tree_graph( + py: Python, + order: u32, + weights: Option>, + multigraph: bool, +) -> PyResult { + let mut graph = StableUnGraph::::default(); + + let num_nodes = usize::pow(2, order); + + let nodes: Vec = match weights { + Some(weights) => { + let mut node_list: Vec = Vec::new(); + + let mut node_count = num_nodes; + + if weights.len() > num_nodes { + return Err(PyIndexError::new_err( + "weights should be <= 2**order", + )); + } + + for weight in weights { + let index = graph.add_node(weight); + node_list.push(index); + node_count -= 1; + } + + for _i in 0..node_count { + let index = graph.add_node(py.None()); + node_list.push(index); + } + + node_list + } + + None => (0..num_nodes).map(|_| graph.add_node(py.None())).collect(), + }; + + let mut n = 1; + + for _ in 0..order { + let edges: Vec<(NodeIndex, NodeIndex)> = graph + .edge_references() + .map(|e| (e.source(), e.target())) + .collect(); + + for (source, target) in edges { + let source_index = source.index(); + let target_index = target.index(); + + graph.add_edge( + nodes[source_index + n], + nodes[target_index + n], + py.None(), + ); + } + + graph.add_edge(nodes[0], nodes[n], py.None()); + + n *= 2; + } + + Ok(graph::PyGraph { + graph, + node_removed: false, + multigraph, + }) +} + +/// Generate an undirected binomial tree of order n recursively. +/// The edges propagate towards right and bottom direction if ``bidirectional`` is ``false`` +/// +/// :param int order: Order of the binomial tree. +/// :param list weights: A list of node weights. If the number of weights is +/// less than 2**order extra nodes with None will be appended. +/// :param bidirectional: A parameter to indicate if edges should exist in +/// both directions between nodes +/// +/// :returns: A directed binomial tree with 2^n vertices and 2^n - 1 edges. +/// :rtype: PyDiGraph +/// :raises IndexError: If the lenght of ``weights`` is greater that 2^n +/// +/// .. jupyter-execute:: +/// +/// import os +/// import tempfile +/// +/// import pydot +/// from PIL import Image +/// +/// import retworkx.generators +/// +/// graph = retworkx.generators.directed_binomial_tree_graph(4) +/// dot_str = graph.to_dot( +/// lambda node: dict( +/// color='black', fillcolor='lightblue', style='filled')) +/// dot = pydot.graph_from_dot_data(dot_str)[0] +/// +/// with tempfile.TemporaryDirectory() as tmpdirname: +/// tmp_path = os.path.join(tmpdirname, 'dag.png') +/// dot.write_png(tmp_path) +/// image = Image.open(tmp_path) +/// os.remove(tmp_path) +/// image +/// +#[pyfunction(bidirectional = "false")] +#[text_signature = "(order, /, weights=None, bidirectional=False)"] +pub fn directed_binomial_tree_graph( + py: Python, + order: u32, + weights: Option>, + bidirectional: bool, +) -> PyResult { + let mut graph = StableDiGraph::::default(); + + let num_nodes = usize::pow(2, order); + + let nodes: Vec = match weights { + Some(weights) => { + let mut node_list: Vec = Vec::new(); + let mut node_count = num_nodes; + + if weights.len() > num_nodes { + return Err(PyIndexError::new_err( + "weights should be <= 2**order", + )); + } + + for weight in weights { + let index = graph.add_node(weight); + node_list.push(index); + node_count -= 1; + } + + for _i in 0..node_count { + let index = graph.add_node(py.None()); + node_list.push(index); + } + + node_list + } + + None => (0..num_nodes).map(|_| graph.add_node(py.None())).collect(), + }; + + let mut n = 1; + + for _ in 0..order { + let edges: Vec<(NodeIndex, NodeIndex)> = graph + .edge_references() + .map(|e| (e.source(), e.target())) + .collect(); + + for (source, target) in edges { + let source_index = source.index(); + let target_index = target.index(); + + if graph + .find_edge(nodes[source_index + n], nodes[target_index + n]) + .is_none() + { + graph.add_edge( + nodes[source_index + n], + nodes[target_index + n], + py.None(), + ); + } + + if bidirectional + && graph + .find_edge(nodes[target_index + n], nodes[source_index + n]) + .is_none() + { + graph.add_edge( + nodes[target_index + n], + nodes[source_index + n], + py.None(), + ); + } + } + + if graph.find_edge(nodes[0], nodes[n]).is_none() { + graph.add_edge(nodes[0], nodes[n], py.None()); + } + + if bidirectional && graph.find_edge(nodes[n], nodes[0]).is_none() { + graph.add_edge(nodes[n], nodes[0], py.None()); + } + + n *= 2; + } + + Ok(digraph::PyDiGraph { + graph, + node_removed: false, + check_cycle: false, + cycle_state: algo::DfsSpace::default(), + multigraph: true, + }) +} + /// Generate an undirected hexagonal lattice graph. /// /// :param int rows: The number of rows to generate the graph with. @@ -1205,6 +1447,8 @@ pub fn generators(_py: Python, m: &PyModule) -> 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!(binomial_tree_graph))?; + m.add_wrapped(wrap_pyfunction!(directed_binomial_tree_graph))?; m.add_wrapped(wrap_pyfunction!(hexagonal_lattice_graph))?; m.add_wrapped(wrap_pyfunction!(directed_hexagonal_lattice_graph))?; Ok(()) diff --git a/tests/generators/test_binomial_tree.py b/tests/generators/test_binomial_tree.py new file mode 100644 index 000000000..02d86369e --- /dev/null +++ b/tests/generators/test_binomial_tree.py @@ -0,0 +1,201 @@ +# 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. + +import unittest + +import retworkx + + +class TestBinomialTreeGraph(unittest.TestCase): + def test_binomial_tree_graph(self): + expected_edges = { + 0: [], + 1: [(0, 1)], + 2: [(0, 1), (2, 3), (0, 2)], + 3: [(0, 1), (2, 3), (0, 2), (4, 5), (6, 7), (4, 6), (0, 4)], + 4: [ + (0, 1), + (2, 3), + (0, 2), + (4, 5), + (6, 7), + (4, 6), + (0, 4), + (8, 9), + (10, 11), + (8, 10), + (12, 13), + (14, 15), + (12, 14), + (8, 12), + (0, 8), + ], + } + for n in range(5): + with self.subTest(n=n): + graph = retworkx.generators.binomial_tree_graph(n) + self.assertEqual(len(graph), 2 ** n) + self.assertEqual(len(graph.edges()), 2 ** n - 1) + self.assertEqual(list(graph.edge_list()), expected_edges[n]) + + def test_binomial_tree_graph_weights(self): + graph = retworkx.generators.binomial_tree_graph( + 2, weights=list(range(4)) + ) + expected_edges = [(0, 1), (2, 3), (0, 2)] + self.assertEqual(len(graph), 4) + self.assertEqual([x for x in range(4)], graph.nodes()) + self.assertEqual(len(graph.edges()), 3) + self.assertEqual(list(graph.edge_list()), expected_edges) + + def test_binomial_tree_graph_weight_less_nodes(self): + graph = retworkx.generators.binomial_tree_graph( + 2, weights=list(range(2)) + ) + self.assertEqual(len(graph), 4) + expected_weights = [x for x in range(2)] + expected_weights.extend([None, None]) + self.assertEqual(expected_weights, graph.nodes()) + self.assertEqual(len(graph.edges()), 3) + + def test_binomial_tree_graph_weights_greater_nodes(self): + with self.assertRaises(IndexError): + retworkx.generators.binomial_tree_graph(2, weights=list(range(7))) + + def test_binomial_tree_no_order(self): + with self.assertRaises(TypeError): + retworkx.generators.binomial_tree_graph(weights=list(range(4))) + + def test_directed_binomial_tree_graph(self): + expected_edges = { + 0: [], + 1: [(0, 1)], + 2: [(0, 1), (2, 3), (0, 2)], + 3: [(0, 1), (2, 3), (0, 2), (4, 5), (6, 7), (4, 6), (0, 4)], + 4: [ + (0, 1), + (2, 3), + (0, 2), + (4, 5), + (6, 7), + (4, 6), + (0, 4), + (8, 9), + (10, 11), + (8, 10), + (12, 13), + (14, 15), + (12, 14), + (8, 12), + (0, 8), + ], + } + + for n in range(5): + with self.subTest(n=n): + graph = retworkx.generators.directed_binomial_tree_graph(n) + self.assertEqual(len(graph), 2 ** n) + self.assertEqual(len(graph.edges()), 2 ** n - 1) + self.assertEqual(list(graph.edge_list()), expected_edges[n]) + + def test_directed_binomial_tree_graph_weights(self): + graph = retworkx.generators.directed_binomial_tree_graph( + 2, weights=list(range(4)) + ) + self.assertEqual(len(graph), 4) + self.assertEqual([x for x in range(4)], graph.nodes()) + self.assertEqual(len(graph.edges()), 3) + + def test_directed_binomial_tree_graph_weight_less_nodes(self): + graph = retworkx.generators.directed_binomial_tree_graph( + 2, weights=list(range(2)) + ) + self.assertEqual(len(graph), 4) + expected_weights = [x for x in range(2)] + expected_weights.extend([None, None]) + self.assertEqual(expected_weights, graph.nodes()) + self.assertEqual(len(graph.edges()), 3) + + def test_directed_binomial_tree_graph_weights_greater_nodes(self): + with self.assertRaises(IndexError): + retworkx.generators.directed_binomial_tree_graph( + 2, weights=list(range(7)) + ) + + def test_directed_binomial_tree_no_order(self): + with self.assertRaises(TypeError): + retworkx.generators.directed_binomial_tree_graph( + weights=list(range(4)) + ) + + def test_directed_binomial_tree_graph_bidirectional(self): + expected_edges = { + 0: [], + 1: [(0, 1), (1, 0)], + 2: [(0, 1), (1, 0), (2, 3), (3, 2), (0, 2), (2, 0)], + 3: [ + (0, 1), + (1, 0), + (2, 3), + (3, 2), + (0, 2), + (2, 0), + (4, 5), + (5, 4), + (6, 7), + (7, 6), + (4, 6), + (6, 4), + (0, 4), + (4, 0), + ], + 4: [ + (0, 1), + (1, 0), + (2, 3), + (3, 2), + (0, 2), + (2, 0), + (4, 5), + (5, 4), + (6, 7), + (7, 6), + (4, 6), + (6, 4), + (0, 4), + (4, 0), + (8, 9), + (9, 8), + (10, 11), + (11, 10), + (8, 10), + (10, 8), + (12, 13), + (13, 12), + (14, 15), + (15, 14), + (12, 14), + (14, 12), + (8, 12), + (12, 8), + (0, 8), + (8, 0), + ], + } + for n in range(5): + with self.subTest(n=n): + graph = retworkx.generators.directed_binomial_tree_graph( + n, bidirectional=True + ) + self.assertEqual(len(graph), 2 ** n) + self.assertEqual(len(graph.edges()), 2 * (2 ** n - 1)) + self.assertEqual(list(graph.edge_list()), expected_edges[n])