feat(graph): add bfs

src/power_grid_model_ds/_core/model/graphs/models/base.py

@@ -4,7 +4,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from collections import Counter
-from collections.abc import Generator
+from collections.abc import Generator, Sequence
 from contextlib import contextmanager
 from itertools import combinations
 from typing import TYPE_CHECKING
@@ -382,6 +382,25 @@ def get_downstream_nodes(self, node_id: int, start_node_ids: list[int], inclusiv
 
         return self.get_connected(node_id, [upstream_node], inclusive)
 
+    def bfs(self, source: int | Sequence[int]) -> list[tuple[int, int | None]]:
+        """Breadth first search from the source(s).
+
+        Args:
+            source(int | Sequence[int]): the source(s) the start the breadth first search from.
+
+        Returns:
+            list[tuple[int, int | None]]: the (node, parent) by the order in which we found the nodes.
+                The parent is None for a source that has not been found via another source yet."""
+        internal_sources = self._externals_to_internals([source] if isinstance(source, int) else source)
+        internal_result = self._bfs(internal_sources)
+        return [
+            (
+                self.internal_to_external(node),
+                None if parent is None else self.internal_to_external(parent),
+            )
+            for node, parent in internal_result
+        ]
+
     def find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph.
         Returns:
@@ -419,7 +438,7 @@ def _internals_to_externals(self, internal_nodes: list[int]) -> list[int]:
         """Convert a list of internal node ids to external node ids"""
         return [self.internal_to_external(node_id) for node_id in internal_nodes]
 
-    def _externals_to_internals(self, external_nodes: list[int] | NDArray) -> list[int]:
+    def _externals_to_internals(self, external_nodes: Sequence[int] | NDArray) -> list[int]:
         """Convert a list of external nodes to internal nodes"""
         return [self.external_to_internal(node_id) for node_id in external_nodes]
 
@@ -526,6 +545,9 @@ def _get_all_paths(self, source, target) -> list[list[int]]: ...
     @abstractmethod
     def _get_components(self) -> list[list[int]]: ...
 
+    @abstractmethod
+    def _bfs(self, source: list[int]) -> list[tuple[int, int | None]]: ...
+
     @abstractmethod
     def _find_fundamental_cycles(self) -> list[list[int]]: ...
 

src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py

@@ -88,6 +88,11 @@ def _delete_branch(self, from_node_id: int, to_node_id: int) -> None:
         except NoEdgeBetweenNodes as error:
             raise MissingBranchError(f"No edge between (internal) nodes {from_node_id} and {to_node_id}") from error
 
+    def _bfs(self, source: list[int]) -> list[tuple[int, int | None]]:
+        visitor = _BfsParentVisitor()
+        rx.bfs_search(self._graph, source, visitor)
+        return [(node, visitor.parents.get(node)) for node in visitor.nodes]
+
     def _get_shortest_path(self, source: int, target: int) -> tuple[list[int], int]:
         path_mapping = rx.dijkstra_shortest_paths(self._graph, source, target)
 
@@ -135,6 +140,19 @@ def _all_branches(self) -> Generator[tuple[int, int], None, None]:
         return ((source, target) for source, target in self._graph.edge_list())
 
 
+class _BfsParentVisitor(BFSVisitor):
+    def __init__(self):
+        self.nodes = []
+        self.parents = {}
+
+    def discover_vertex(self, v):
+        self.nodes.append(v)
+
+    def tree_edge(self, e):
+        (u, v, _) = e
+        self.parents[v] = u
+
+
 class _NodeVisitor(BFSVisitor):
     def __init__(self, nodes_to_ignore: list[int]):
         self.nodes_to_ignore = nodes_to_ignore

tests/unit/model/graphs/test_graph_model.py

@@ -476,3 +476,26 @@ def test_has_parallel_edges_reversed(self, graph: BaseGraphModel):
         assert not graph.has_parallel_edges()
         graph.add_branch(2, 1)
         assert graph.has_parallel_edges()
+
+
+class TestBfsSearch:
+    @pytest.mark.parametrize(
+        ("source", "expected"),
+        [
+            pytest.param(1, [(1, None), (5, 1), (2, 1), (4, 5), (3, 2)], id="source: 1"),
+            pytest.param({1}, [(1, None), (5, 1), (2, 1), (4, 5), (3, 2)], id="source {1}"),
+            pytest.param([1, 2], [(1, None), (5, 1), (2, 1), (4, 5), (3, 2)], id="source [1,2]"),
+            pytest.param([2, 1], [(2, None), (3, 2), (1, 2), (5, 1), (4, 5)], id="source [1,2]"),
+            pytest.param({}, [], id="empty source"),
+        ],
+    )
+    def test_bfs_int_source(self, graph_with_2_routes, source, expected):
+        assert graph_with_2_routes.bfs(source) == expected
+
+    def test_bfs_non_existing_node(self, graph_with_2_routes):
+        with pytest.raises(MissingNodeError, match="External node id '10' does NOT exist"):
+            assert graph_with_2_routes.bfs(10)
+
+    def test_second_source_in_different_component(self, graph_with_2_routes):
+        graph_with_2_routes.add_node(8)
+        assert graph_with_2_routes.bfs([1, 8]) == [(1, None), (5, 1), (2, 1), (4, 5), (3, 2), (8, None)]