merge main

Signed-off-by: Thijs Baaijen <[email protected]>

Author: Thijss

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

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: MPL-2.0
 
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from typing import Generator
 
 import numpy as np
 from numpy._typing import NDArray
@@ -34,6 +36,14 @@ def nr_nodes(self):
     def nr_branches(self):
         """Returns the number of branches in the graph"""
 
+    @property
+    def all_branches(self) -> Generator[tuple[int, int], None, None]:
+        """Returns all branches in the graph."""
+        return (
+            (self.internal_to_external(source), self.internal_to_external(target))
+            for source, target in self._all_branches()
+        )
+
     @abstractmethod
     def external_to_internal(self, ext_node_id: int) -> int:
         """Convert external node id to internal node id (internal)
@@ -63,6 +73,14 @@ def has_node(self, node_id: int) -> bool:
 
         return self._has_node(node_id=internal_node_id)
 
+    def in_branches(self, node_id: int) -> Generator[tuple[int, int], None, None]:
+        """Return all branches that have the node as an endpoint."""
+        int_node_id = self.external_to_internal(node_id)
+        internal_edges = self._in_branches(int_node_id=int_node_id)
+        return (
+            (self.internal_to_external(source), self.internal_to_external(target)) for source, target in internal_edges
+        )
+
     def add_node(self, ext_node_id: int, raise_on_fail: bool = True) -> None:
         """Add a node to the graph."""
         if self.has_node(ext_node_id):
@@ -164,6 +182,28 @@ def delete_branch3_array(self, branch3_array: Branch3Array, raise_on_fail: bool
             branches = _get_branch3_branches(branch3)
             self.delete_branch_array(branches, raise_on_fail=raise_on_fail)
 
+    @contextmanager
+    def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
+        """Context manager that temporarily removes nodes and their branches from the graph.
+        Example:
+            >>> with graph.tmp_remove_nodes([1, 2, 3]):
+            >>>    assert not graph.has_node(1)
+            >>> assert graph.has_node(1)
+        In practice, this is useful when you want to e.g. calculate the shortest path between two nodes without
+        considering certain nodes.
+        """
+        edge_list = []
+        for node in nodes:
+            edge_list += list(self.in_branches(node))
+            self.delete_node(node)
+
+        yield
+
+        for node in nodes:
+            self.add_node(node)
+        for source, target in edge_list:
+            self.add_branch(source, target)
+
     def get_shortest_path(self, ext_start_node_id: int, ext_end_node_id: int) -> tuple[list[int], int]:
         """Calculate the shortest path between two nodes
 
@@ -235,8 +275,49 @@ def get_connected(
             nodes_to_ignore=self._externals_to_internals(nodes_to_ignore),
             inclusive=inclusive,
         )
+
         return self._internals_to_externals(nodes)
 
+    def find_first_connected(self, node_id: int, candidate_node_ids: list[int]) -> int:
+        """Find the first connected node to the node_id from the candidate_node_ids
+
+        Note:
+            If multiple candidate nodes are connected to the node, the first one found is returned.
+            There is no guarantee that the same candidate node will be returned each time.
+
+        Raises:
+            MissingNodeError: if no connected node is found
+            ValueError: if the node_id is in candidate_node_ids
+        """
+        internal_node_id = self.external_to_internal(node_id)
+        internal_candidates = self._externals_to_internals(candidate_node_ids)
+        if internal_node_id in internal_candidates:
+            raise ValueError("node_id cannot be in candidate_node_ids")
+        return self.internal_to_external(self._find_first_connected(internal_node_id, internal_candidates))
+
+    def get_downstream_nodes(self, node_id: int, start_node_ids: list[int], inclusive: bool = False) -> list[int]:
+        """Find all nodes downstream of the node_id with respect to the start_node_ids
+
+        Example:
+            given this graph: [1] - [2] - [3] - [4]
+            >>> graph.get_downstream_nodes(2, [1]) == [3, 4]
+            >>> graph.get_downstream_nodes(2, [1], inclusive=True) == [2, 3, 4]
+
+        args:
+            node_id: node id to start the search from
+            start_node_ids: list of node ids considered 'above' the node_id
+            inclusive: whether to include the given node id in the result
+        returns:
+            list of node ids sorted by distance, downstream of to the node id
+        """
+        connected_node = self.find_first_connected(node_id, start_node_ids)
+        path, _ = self.get_shortest_path(node_id, connected_node)
+        _, upstream_node, *_ = (
+            path  # path is at least 2 elements long or find_first_connected would have raised an error
+        )
+
+        return self.get_connected(node_id, [upstream_node], inclusive)
+
     def find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph.
         Returns:
@@ -270,9 +351,15 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
             return branch.is_active.item()
         return True
 
+    @abstractmethod
+    def _in_branches(self, int_node_id: int) -> Generator[tuple[int, int], None, None]: ...
+
     @abstractmethod
     def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bool = False) -> list[int]: ...
 
+    @abstractmethod
+    def _find_first_connected(self, node_id: int, candidate_node_ids: list[int]) -> int: ...
+
     @abstractmethod
     def _has_branch(self, from_node_id, to_node_id) -> bool: ...
 
@@ -307,6 +394,9 @@ def _get_components(self, substation_nodes: list[int]) -> list[list[int]]: ...
     @abstractmethod
     def _find_fundamental_cycles(self) -> list[list[int]]: ...
 
+    @abstractmethod
+    def _all_branches(self) -> Generator[tuple[int, int], None, None]: ...
+
 
 def _get_branch3_branches(branch3: Branch3Array) -> BranchArray:
     node_1 = branch3.node_1.item()

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

@@ -3,10 +3,11 @@
 # SPDX-License-Identifier: MPL-2.0
 
 import logging
+from typing import Generator
 
 import rustworkx as rx
 from rustworkx import NoEdgeBetweenNodes
-from rustworkx.visit import BFSVisitor, PruneSearch
+from rustworkx.visit import BFSVisitor, PruneSearch, StopSearch
 
 from power_grid_model_ds._core.model.graphs.errors import MissingBranchError, MissingNodeError, NoPathBetweenNodes
 from power_grid_model_ds._core.model.graphs.models._rustworkx_search import find_fundamental_cycles_rustworkx
@@ -99,6 +100,16 @@ def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bo
 
         return connected_nodes
 
+    def _in_branches(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
+        return ((source, target) for source, target, _ in self._graph.in_edges(int_node_id))
+
+    def _find_first_connected(self, node_id: int, candidate_node_ids: list[int]) -> int:
+        visitor = _NodeFinder(candidate_nodes=candidate_node_ids)
+        rx.bfs_search(self._graph, [node_id], visitor)
+        if visitor.found_node is None:
+            raise MissingNodeError(f"node {node_id} is not connected to any of the candidate nodes")
+        return visitor.found_node
+
     def _find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph using Rustworkx.
 
@@ -107,6 +118,9 @@ def _find_fundamental_cycles(self) -> list[list[int]]:
         """
         return find_fundamental_cycles_rustworkx(self._graph)
 
+    def _all_branches(self) -> Generator[tuple[int, int], None, None]:
+        return ((source, target) for source, target in self._graph.edge_list())
+
 
 class _NodeVisitor(BFSVisitor):
     def __init__(self, nodes_to_ignore: list[int]):
@@ -117,3 +131,16 @@ def discover_vertex(self, v):
         if v in self.nodes_to_ignore:
             raise PruneSearch
         self.nodes.append(v)
+
+
+class _NodeFinder(BFSVisitor):
+    """Visitor that stops the search when a candidate node is found"""
+
+    def __init__(self, candidate_nodes: list[int]):
+        self.candidate_nodes = candidate_nodes
+        self.found_node: int | None = None
+
+    def discover_vertex(self, v):
+        if v in self.candidate_nodes:
+            self.found_node = v
+            raise StopSearch

src/power_grid_model_ds/_core/model/grids/base.py

@@ -331,6 +331,7 @@ def get_nearest_substation_node(self, node_id: int):
 
     def get_downstream_nodes(self, node_id: int, inclusive: bool = False):
         """Get the downstream nodes from a node.
+        Assuming each node has a single feeding substation and the grid is radial
 
         Example:
             given this graph: [1] - [2] - [3] - [4], with 1 being a substation node
@@ -349,15 +350,14 @@ def get_downstream_nodes(self, node_id: int, inclusive: bool = False):
         Returns:
             list[int]: The downstream nodes.
         """
-        substation_node_id = self.get_nearest_substation_node(node_id).id.item()
+        substation_nodes = self.node.filter(node_type=NodeType.SUBSTATION_NODE.value)
 
-        if node_id == substation_node_id:
+        if node_id in substation_nodes.id:
             raise NotImplementedError("get_downstream_nodes is not implemented for substation nodes!")
 
-        path_to_substation, _ = self.graphs.active_graph.get_shortest_path(node_id, substation_node_id)
-        upstream_node = path_to_substation[1]
-
-        return self.graphs.active_graph.get_connected(node_id, nodes_to_ignore=[upstream_node], inclusive=inclusive)
+        return self.graphs.active_graph.get_downstream_nodes(
+            node_id=node_id, start_node_ids=list(substation_nodes.id), inclusive=inclusive
+        )
 
     def cache(self, cache_dir: Path, cache_name: str, compress: bool = True):
         """Cache Grid to a folder

tests/unit/model/graphs/test_graph_model.py

@@ -1,6 +1,11 @@
 # SPDX-FileCopyrightText: Contributors to the Power Grid Model project <[email protected]>
 #
 # SPDX-License-Identifier: MPL-2.0
+
+"""Grid tests"""
+
+from collections import Counter
+
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
@@ -10,7 +15,7 @@
 # pylint: disable=missing-function-docstring,missing-class-docstring
 
 
-class TestGraphModifications:
+class TestBasicGraphFunctions:
     def test_graph_add_node_and_branch(self, graph):
         graph.add_node(1)
         graph.add_node(2)
@@ -23,6 +28,15 @@ def test_graph_add_node_and_branch(self, graph):
         assert 2 == graph.nr_nodes
         assert 1 == graph.nr_branches
 
+    def test_add_invalid_branch(self, graph):
+        graph.add_node(1)
+        graph.add_node(2)
+        graph.add_branch(1, 2)
+        assert graph.has_branch(1, 2)
+
+        with pytest.raises(MissingNodeError):
+            graph.add_branch(1, 3)
+
     def test_has_node(self, graph):
         graph.add_node(1)
         assert graph.has_node(1)
@@ -37,14 +51,21 @@ def test_has_branch(self, graph):
         assert graph.has_branch(2, 1)  # reversed should work too
         assert not graph.has_branch(1, 3)
 
-    def test_add_invalid_branch(self, graph):
+    def test_graph_all_branches(self, graph):
         graph.add_node(1)
         graph.add_node(2)
         graph.add_branch(1, 2)
-        assert graph.has_branch(1, 2)
 
-        with pytest.raises(MissingNodeError):
-            graph.add_branch(1, 3)
+        assert [(1, 2)] == list(graph.all_branches)
+
+    def test_graph_all_branches_parallel(self, graph):
+        graph.add_node(1)
+        graph.add_node(2)
+        graph.add_branch(1, 2)
+        graph.add_branch(1, 2)
+        graph.add_branch(2, 1)
+
+        assert [(1, 2), (1, 2), (2, 1)] == list(graph.all_branches)
 
     def test_delete_invalid_node_without_error(self, graph):
         graph.delete_node(3, raise_on_fail=False)
@@ -104,6 +125,62 @@ def test_internal_ids_after_node_deletion(self, graph):
         assert graph._has_node(internal_id_0)
         assert graph._has_node(internal_id_2)
 
+    def test_graph_in_branches(self, graph):
+        graph.add_node(1)
+        graph.add_node(2)
+        graph.add_branch(1, 2)
+        graph.add_branch(1, 2)
+        graph.add_branch(2, 1)
+
+        assert [(2, 1), (2, 1), (2, 1)] == list(graph.in_branches(1))
+        assert [(1, 2), (1, 2), (1, 2)] == list(graph.in_branches(2))
+
+
+def test_tmp_remove_nodes(graph_with_2_routes) -> None:
+    graph = graph_with_2_routes
+
+    assert graph.nr_branches == 4
+
+    # add parallel branches to test whether they are restored correctly
+    graph.add_branch(1, 5)
+    graph.add_branch(5, 1)
+
+    assert graph.nr_nodes == 5
+    assert graph.nr_branches == 6
+
+    before_sets = [frozenset(branch) for branch in graph.all_branches]
+    counter_before = Counter(before_sets)
+
+    with graph.tmp_remove_nodes([1, 2]):
+        assert graph.nr_nodes == 3
+        assert list(graph.all_branches) == [(5, 4)]
+
+    assert graph.nr_nodes == 5
+    assert graph.nr_branches == 6
+
+    after_sets = [frozenset(branch) for branch in graph.all_branches]
+    counter_after = Counter(after_sets)
+    assert counter_before == counter_after
+
+
+def test_get_components(graph_with_2_routes):
+    graph = graph_with_2_routes
+    graph.add_node(99)
+    graph.add_branch(1, 99)
+    substation_nodes = np.array([1])
+
+    components = graph.get_components(substation_nodes=substation_nodes)
+
+    assert len(components) == 3
+    assert set(components[0]) == {2, 3}
+    assert set(components[1]) == {4, 5}
+    assert set(components[2]) == {99}
+
+
+def test_from_arrays(basic_grid):
+    new_graph = basic_grid.graphs.complete_graph.__class__.from_arrays(basic_grid)
+    assert_array_equal(new_graph.external_ids, basic_grid.node.id)
+
 
 class TestPathMethods:
     def test_get_shortest_path(self, graph_with_2_routes):
@@ -248,20 +325,18 @@ def test_get_connected_ignore_multiple_nodes(self, graph_with_2_routes):
         assert {5} == set(connected_nodes)
 
 
-def test_get_components(graph_with_2_routes):
-    graph = graph_with_2_routes
-    graph.add_node(99)
-    graph.add_branch(1, 99)
-    substation_nodes = np.array([1])
-
-    components = graph.get_components(substation_nodes=substation_nodes)
-
-    assert len(components) == 3
-    assert set(components[0]) == {2, 3}
-    assert set(components[1]) == {4, 5}
-    assert set(components[2]) == {99}
+class TestFindFirstConnected:
+    def test_find_first_connected(self, graph_with_2_routes):
+        graph = graph_with_2_routes
+        assert 2 == graph.find_first_connected(1, candidate_node_ids=[2, 3, 4])
 
+    def test_find_first_connected_same_node(self, graph_with_2_routes):
+        graph = graph_with_2_routes
+        with pytest.raises(ValueError):
+            graph.find_first_connected(1, candidate_node_ids=[1, 3, 5])
 
-def test_from_arrays(basic_grid):
-    new_graph = basic_grid.graphs.complete_graph.__class__.from_arrays(basic_grid)
-    assert_array_equal(new_graph.external_ids, basic_grid.node.id)
+    def test_find_first_connected_no_match(self, graph_with_2_routes):
+        graph = graph_with_2_routes
+        graph.add_node(99)
+        with pytest.raises(MissingNodeError):
+            graph.find_first_connected(1, candidate_node_ids=[99])