Skip to content

Commit b7f8afb

Browse files
committed
merge all-branches branch
Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
2 parents a94b8f2 + f237fe2 commit b7f8afb

File tree

3 files changed

+34
-0
lines changed

3 files changed

+34
-0
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,13 @@ def nr_nodes(self):
3636
def nr_branches(self):
3737
"""Returns the number of branches in the graph"""
3838

39+
@property
40+
@abstractmethod
41+
def all_branches(self) -> list[frozenset[int]]:
42+
"""Returns all branches in the graph as a list of node pairs (frozensets).
43+
Warning: Depending on graph engine, performance could be slow for large graphs
44+
"""
45+
3946
@abstractmethod
4047
def external_to_internal(self, ext_node_id: int) -> int:
4148
"""Convert external node id to internal node id (internal)

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,15 @@ def nr_nodes(self):
3232
def nr_branches(self):
3333
return self._graph.num_edges()
3434

35+
@property
36+
def all_branches(self) -> list[frozenset[int]]:
37+
internal_branches = ((source, target) for source, target in self._graph.edge_list())
38+
external_branches = [
39+
frozenset([self.internal_to_external(source), self.internal_to_external(target)])
40+
for source, target in internal_branches
41+
]
42+
return external_branches
43+
3544
@property
3645
def external_ids(self) -> list[int]:
3746
return list(self._external_to_internal.keys())

tests/unit/model/graphs/test_graph_model.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,24 @@ def test_graph_has_branch(graph):
3737
assert not graph.has_branch(1, 3)
3838

3939

40+
def test_graph_all_branches(graph):
41+
graph.add_node(1)
42+
graph.add_node(2)
43+
graph.add_branch(1, 2)
44+
45+
assert [{1, 2}] == graph.all_branches
46+
47+
48+
def test_graph_all_branches_parallel(graph):
49+
graph.add_node(1)
50+
graph.add_node(2)
51+
graph.add_branch(1, 2)
52+
graph.add_branch(1, 2)
53+
graph.add_branch(2, 1)
54+
55+
assert [{1, 2}, {1, 2}, {1, 2}] == graph.all_branches
56+
57+
4058
def test_graph_delete_branch(graph):
4159
"""Test whether a branch is deleted correctly"""
4260
graph.add_node(1)

0 commit comments

Comments
 (0)