Add tmp-remove-nodes method

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>

Author: Thijss

VERSION

@@ -1 +1 @@
-1.0
+1.1

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
@@ -164,6 +166,31 @@ def delete_branch3_array(self, branch_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:
+            internal_node = self.external_to_internal(node)
+            node_edges = [
+                (self.internal_to_external(source), self.internal_to_external(target))
+                for source, target in self._in_edges(internal_node)
+            ]
+            edge_list += node_edges
+            self._delete_node(internal_node)
+        yield edge_list
+        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
 
@@ -270,6 +297,13 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
             return branch.is_active.item()
         return True
 
+    @abstractmethod
+    def _in_edges(self, internal_node: int) -> list[tuple[int, int]]:
+        """Return all edges a node occurs in.
+
+        Return a list of tuples with the source and target node id. These are internal node ids.
+        """
+
     @abstractmethod
     def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bool = False) -> list[int]: ...
 

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

@@ -99,6 +99,9 @@ def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bo
 
         return connected_nodes
 
+    def _in_edges(self, internal_node: int) -> list[tuple[int, int]]:
+        return [(source, target) for source, target, _ in self._graph.in_edges(internal_node)]
+
     def _find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph using Rustworkx.
 

tests/unit/model/graphs/test_graph_model.py

@@ -27,6 +27,16 @@ def test_graph_initialize(graph):
     assert 1 == graph.nr_branches
 
 
+def test_graph_has_branch(graph):
+    graph.add_node(1)
+    graph.add_node(2)
+    graph.add_branch(1, 2)
+
+    assert graph.has_branch(1, 2)
+    assert graph.has_branch(2, 1)  # reversed should work too
+    assert not graph.has_branch(1, 3)
+
+
 def test_graph_delete_branch(graph):
     """Test whether a branch is deleted correctly"""
     graph.add_node(1)