From f237fe2f227dffcd3fafc473e317b3af29828130 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+thijss@users.noreply.github.com>
Date: Thu, 23 Jan 2025 16:06:16 +0100
Subject: [PATCH 01/10] Add .all_branches property to BaseGraphModel

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 VERSION                                        |  2 +-
 .../_core/model/graphs/models/base.py          |  7 +++++++
 .../_core/model/graphs/models/rustworkx.py     |  9 +++++++++
 tests/unit/model/graphs/test_graph_model.py    | 18 ++++++++++++++++++
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/VERSION b/VERSION
index 9f8e9b6..b123147 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-1.0
\ No newline at end of file
+1.1
\ No newline at end of file
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index 2a4a577..a44b01a 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -34,6 +34,13 @@ def nr_nodes(self):
     def nr_branches(self):
         """Returns the number of branches in the graph"""
 
+    @property
+    @abstractmethod
+    def all_branches(self) -> list[frozenset[int]]:
+        """Returns all branches in the graph as a list of node pairs (frozensets).
+        Warning: Depending on graph engine, performance could be slow for large graphs
+        """
+
     @abstractmethod
     def external_to_internal(self, ext_node_id: int) -> int:
         """Convert external node id to internal node id (internal)
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index 6e1b7cd..012c487 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
@@ -32,6 +32,15 @@ def nr_nodes(self):
     def nr_branches(self):
         return self._graph.num_edges()
 
+    @property
+    def all_branches(self) -> list[frozenset[int]]:
+        internal_branches = ((source, target) for source, target in self._graph.edge_list())
+        external_branches = [
+            frozenset([self.internal_to_external(source), self.internal_to_external(target)])
+            for source, target in internal_branches
+        ]
+        return external_branches
+
     @property
     def external_ids(self) -> list[int]:
         return list(self._external_to_internal.keys())
diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index 628c77b..f259864 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/tests/unit/model/graphs/test_graph_model.py
@@ -37,6 +37,24 @@ def test_graph_has_branch(graph):
     assert not graph.has_branch(1, 3)
 
 
+def test_graph_all_branches(graph):
+    graph.add_node(1)
+    graph.add_node(2)
+    graph.add_branch(1, 2)
+
+    assert [{1, 2}] == graph.all_branches
+
+
+def test_graph_all_branches_parallel(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}, {1, 2}] == graph.all_branches
+
+
 def test_graph_delete_branch(graph):
     """Test whether a branch is deleted correctly"""
     graph.add_node(1)

From 1154f96e1dd2772fd4b7f751333a27a121732cf9 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Thu, 23 Jan 2025 15:35:18 +0100
Subject: [PATCH 02/10] Add tmp-remove-nodes method

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 VERSION                                       |  2 +-
 .../_core/model/graphs/models/base.py         | 34 +++++++++++++++++++
 .../_core/model/graphs/models/rustworkx.py    |  3 ++
 tests/unit/model/graphs/test_graph_model.py   | 10 ++++++
 4 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/VERSION b/VERSION
index 9f8e9b6..b123147 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-1.0
\ No newline at end of file
+1.1
\ No newline at end of file
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index 2a4a577..33f194b 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/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]: ...
 
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index 6e1b7cd..12c638b 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/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.
 
diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index cea5477..628c77b 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/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)

From bcd687008e56edda05e5c2fa1dcbc4cabeabffc5 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+thijss@users.noreply.github.com>
Date: Thu, 23 Jan 2025 16:06:16 +0100
Subject: [PATCH 03/10] Add .all_branches property to BaseGraphModel

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 .../_core/model/graphs/models/base.py          |  7 +++++++
 .../_core/model/graphs/models/rustworkx.py     |  9 +++++++++
 tests/unit/model/graphs/test_graph_model.py    | 18 ++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index 33f194b..9b37449 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -36,6 +36,13 @@ def nr_nodes(self):
     def nr_branches(self):
         """Returns the number of branches in the graph"""
 
+    @property
+    @abstractmethod
+    def all_branches(self) -> list[frozenset[int]]:
+        """Returns all branches in the graph as a list of node pairs (frozensets).
+        Warning: Depending on graph engine, performance could be slow for large graphs
+        """
+
     @abstractmethod
     def external_to_internal(self, ext_node_id: int) -> int:
         """Convert external node id to internal node id (internal)
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index 12c638b..bf4df27 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
@@ -32,6 +32,15 @@ def nr_nodes(self):
     def nr_branches(self):
         return self._graph.num_edges()
 
+    @property
+    def all_branches(self) -> list[frozenset[int]]:
+        internal_branches = ((source, target) for source, target in self._graph.edge_list())
+        external_branches = [
+            frozenset([self.internal_to_external(source), self.internal_to_external(target)])
+            for source, target in internal_branches
+        ]
+        return external_branches
+
     @property
     def external_ids(self) -> list[int]:
         return list(self._external_to_internal.keys())
diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index 628c77b..f259864 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/tests/unit/model/graphs/test_graph_model.py
@@ -37,6 +37,24 @@ def test_graph_has_branch(graph):
     assert not graph.has_branch(1, 3)
 
 
+def test_graph_all_branches(graph):
+    graph.add_node(1)
+    graph.add_node(2)
+    graph.add_branch(1, 2)
+
+    assert [{1, 2}] == graph.all_branches
+
+
+def test_graph_all_branches_parallel(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}, {1, 2}] == graph.all_branches
+
+
 def test_graph_delete_branch(graph):
     """Test whether a branch is deleted correctly"""
     graph.add_node(1)

From 45780e70cab14ff590506f3354a698b7f21dbdcb Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Thu, 23 Jan 2025 16:17:56 +0100
Subject: [PATCH 04/10] Add test

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/unit/model/graphs/test_graph_model.py | 25 +++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index f259864..d8e1128 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/tests/unit/model/graphs/test_graph_model.py
@@ -4,6 +4,8 @@
 
 """Grid tests"""
 
+from collections import Counter
+
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
@@ -338,3 +340,26 @@ def test_get_connected_ignore_multiple_nodes(self, graph_with_2_routes):
         connected_nodes = graph.get_connected(node_id=1, nodes_to_ignore=[2, 4])
 
         assert {5} == set(connected_nodes)
+
+
+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
+    counter_before: Counter[frozenset] = Counter(graph.all_branches)
+
+    with graph.tmp_remove_nodes([1, 2]):
+        assert graph.nr_nodes == 3
+        assert graph.all_branches == [{4, 5}]
+
+    assert graph.nr_nodes == 5
+    assert graph.nr_branches == 6
+    counter_after: Counter[frozenset] = Counter(graph.all_branches)
+    assert counter_before == counter_after

From 3daad3f4201cd0867e72301e40b0002147023a80 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Fri, 24 Jan 2025 12:20:51 +0100
Subject: [PATCH 05/10] Apply suggestions from code review

Co-authored-by: Vincent Koppen <vincent.koppen@alliander.com>
Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 .../_core/model/graphs/models/base.py             | 15 ++++++++++-----
 .../_core/model/graphs/models/rustworkx.py        | 13 ++++---------
 tests/unit/model/graphs/test_graph_model.py       |  4 ++--
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index a44b01a..6b36e5a 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 from abc import ABC, abstractmethod
+from typing import Generator
 
 import numpy as np
 from numpy._typing import NDArray
@@ -35,11 +36,12 @@ def nr_branches(self):
         """Returns the number of branches in the graph"""
 
     @property
-    @abstractmethod
-    def all_branches(self) -> list[frozenset[int]]:
-        """Returns all branches in the graph as a list of node pairs (frozensets).
-        Warning: Depending on graph engine, performance could be slow for large graphs
-        """
+    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:
@@ -314,6 +316,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()
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index 012c487..ee2cf46 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 import logging
+from typing import Generator
 
 import rustworkx as rx
 from rustworkx import NoEdgeBetweenNodes
@@ -32,15 +33,6 @@ def nr_nodes(self):
     def nr_branches(self):
         return self._graph.num_edges()
 
-    @property
-    def all_branches(self) -> list[frozenset[int]]:
-        internal_branches = ((source, target) for source, target in self._graph.edge_list())
-        external_branches = [
-            frozenset([self.internal_to_external(source), self.internal_to_external(target)])
-            for source, target in internal_branches
-        ]
-        return external_branches
-
     @property
     def external_ids(self) -> list[int]:
         return list(self._external_to_internal.keys())
@@ -116,6 +108,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]):
diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index f259864..15e250e 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/tests/unit/model/graphs/test_graph_model.py
@@ -42,7 +42,7 @@ def test_graph_all_branches(graph):
     graph.add_node(2)
     graph.add_branch(1, 2)
 
-    assert [{1, 2}] == graph.all_branches
+    assert [(1, 2)] == list(graph.all_branches)
 
 
 def test_graph_all_branches_parallel(graph):
@@ -52,7 +52,7 @@ def test_graph_all_branches_parallel(graph):
     graph.add_branch(1, 2)
     graph.add_branch(2, 1)
 
-    assert [{1, 2}, {1, 2}, {1, 2}] == graph.all_branches
+    assert [(1, 2), (1, 2), (2, 1)] == list(graph.all_branches)
 
 
 def test_graph_delete_branch(graph):

From 4cae2a45bec0d4b74a5c3c64157c3b87ec080ff2 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Tue, 28 Jan 2025 10:25:23 +0100
Subject: [PATCH 06/10] yield nothing

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 src/power_grid_model_ds/_core/model/graphs/models/base.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index 1818f06..9c3f0bd 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from typing import Generator
 
 import numpy as np
@@ -192,7 +193,8 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
             ]
             edge_list += node_edges
             self._delete_node(internal_node)
-        yield edge_list
+        yield
+
         for node in nodes:
             self.add_node(node)
         for source, target in edge_list:

From 4cb80a42c72587ad80c92582b328bdb21b3f9e2e Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Tue, 28 Jan 2025 10:36:01 +0100
Subject: [PATCH 07/10] fix merge issues

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 .../_core/model/graphs/models/rustworkx.py             |  9 ---------
 tests/unit/model/graphs/test_graph_model.py            | 10 +++++++---
 2 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index bfe0c89..397297d 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
@@ -33,15 +33,6 @@ def nr_nodes(self):
     def nr_branches(self):
         return self._graph.num_edges()
 
-    @property
-    def all_branches(self) -> list[frozenset[int]]:
-        internal_branches = ((source, target) for source, target in self._graph.edge_list())
-        external_branches = [
-            frozenset([self.internal_to_external(source), self.internal_to_external(target)])
-            for source, target in internal_branches
-        ]
-        return external_branches
-
     @property
     def external_ids(self) -> list[int]:
         return list(self._external_to_internal.keys())
diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index a05af7f..579f57f 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/tests/unit/model/graphs/test_graph_model.py
@@ -353,13 +353,17 @@ def test_tmp_remove_nodes(graph_with_2_routes) -> None:
 
     assert graph.nr_nodes == 5
     assert graph.nr_branches == 6
-    counter_before: Counter[frozenset] = Counter(graph.all_branches)
+
+    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 graph.all_branches == [{4, 5}]
+        assert list(graph.all_branches) == [(5, 4)]
 
     assert graph.nr_nodes == 5
     assert graph.nr_branches == 6
-    counter_after: Counter[frozenset] = Counter(graph.all_branches)
+
+    after_sets = [frozenset(branch) for branch in graph.all_branches]
+    counter_after = Counter(after_sets)
     assert counter_before == counter_after

From 369498b4b9c92938286e5f9bb2d589e7a51846b7 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Tue, 28 Jan 2025 10:46:06 +0100
Subject: [PATCH 08/10] make .in_edges a public method

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 .../_core/model/graphs/models/base.py         | 20 +++++++++++--------
 .../_core/model/graphs/models/rustworkx.py    |  4 ++--
 tests/unit/model/graphs/test_graph_model.py   | 11 ++++++++++
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index 9c3f0bd..edb4dad 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -73,6 +73,14 @@ def has_node(self, node_id: int) -> bool:
 
         return self._has_node(node_id=internal_node_id)
 
+    def in_edges(self, node_id: int) -> Generator[tuple[int, int], None, None]:
+        """Return all edges a node occurs in."""
+        int_node_id = self.external_to_internal(node_id)
+        internal_edges = self._in_edges(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):
@@ -187,11 +195,7 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
         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
+            edge_list += list(self.in_edges(node))
             self._delete_node(internal_node)
         yield
 
@@ -307,10 +311,10 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
         return True
 
     @abstractmethod
-    def _in_edges(self, internal_node: int) -> list[tuple[int, int]]:
+    def _in_edges(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
         """Return all edges a node occurs in.
-
-        Return a list of tuples with the source and target node id. These are internal node ids.
+        Return a list of tuples with the source and target node id.
+        These are internal node ids.
         """
 
     @abstractmethod
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index 397297d..1c10753 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
@@ -100,8 +100,8 @@ 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 _in_edges(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_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph using Rustworkx.
diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index 579f57f..053c782 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/tests/unit/model/graphs/test_graph_model.py
@@ -57,6 +57,17 @@ def test_graph_all_branches_parallel(graph):
     assert [(1, 2), (1, 2), (2, 1)] == list(graph.all_branches)
 
 
+def test_graph_in_edges(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_edges(1))
+    assert [(1, 2), (1, 2), (1, 2)] == list(graph.in_edges(2))
+
+
 def test_graph_delete_branch(graph):
     """Test whether a branch is deleted correctly"""
     graph.add_node(1)

From e81346f500f09a9bec823919fc32968480b2e9d0 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Tue, 28 Jan 2025 16:43:10 +0100
Subject: [PATCH 09/10] rename method to .in_branches

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 .../_core/model/graphs/models/base.py              | 14 +++++++-------
 .../_core/model/graphs/models/rustworkx.py         |  2 +-
 tests/unit/model/graphs/test_graph_model.py        |  6 +++---
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index edb4dad..aa1c07e 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -73,10 +73,10 @@ def has_node(self, node_id: int) -> bool:
 
         return self._has_node(node_id=internal_node_id)
 
-    def in_edges(self, node_id: int) -> Generator[tuple[int, int], None, None]:
-        """Return all edges a node occurs in."""
+    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_edges(int_node_id=int_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
         )
@@ -194,9 +194,9 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
         """
         edge_list = []
         for node in nodes:
-            internal_node = self.external_to_internal(node)
-            edge_list += list(self.in_edges(node))
-            self._delete_node(internal_node)
+            edge_list += list(self.in_branches(node))
+            self.delete_node(node)
+
         yield
 
         for node in nodes:
@@ -311,7 +311,7 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
         return True
 
     @abstractmethod
-    def _in_edges(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
+    def _in_branches(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
         """Return all edges a node occurs in.
         Return a list of tuples with the source and target node id.
         These are internal node ids.
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index 1c10753..b509391 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
@@ -100,7 +100,7 @@ def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bo
 
         return connected_nodes
 
-    def _in_edges(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
+    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_fundamental_cycles(self) -> list[list[int]]:
diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py
index 053c782..ca9e137 100644
--- a/tests/unit/model/graphs/test_graph_model.py
+++ b/tests/unit/model/graphs/test_graph_model.py
@@ -57,15 +57,15 @@ def test_graph_all_branches_parallel(graph):
     assert [(1, 2), (1, 2), (2, 1)] == list(graph.all_branches)
 
 
-def test_graph_in_edges(graph):
+def test_graph_in_branches(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_edges(1))
-    assert [(1, 2), (1, 2), (1, 2)] == list(graph.in_edges(2))
+    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_graph_delete_branch(graph):

From c223e8325d748c446d6facc4fd7fe5b41c5b529b Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Tue, 28 Jan 2025 17:48:51 +0100
Subject: [PATCH 10/10] remove unnecessary docstring

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 src/power_grid_model_ds/_core/model/graphs/models/base.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index aa1c07e..4f0a7bd 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -311,11 +311,7 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
         return True
 
     @abstractmethod
-    def _in_branches(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
-        """Return all edges a node occurs in.
-        Return a list of tuples with the source and target node id.
-        These are internal node ids.
-        """
+    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]: ...