simplify function

Signed-off-by: Vincent Koppen <vincent.koppen@alliander.com>

Author: vincentkoppen

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

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 import dataclasses
-from collections import defaultdict
+from itertools import pairwise
 from typing import TYPE_CHECKING, Iterator
 
 import numpy as np
@@ -12,6 +12,7 @@
 from power_grid_model_ds._core import fancypy as fp
 from power_grid_model_ds._core.model.arrays import BranchArray
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
+from power_grid_model_ds._core.model.arrays.pgm_arrays import Branch3Array
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 from power_grid_model_ds._core.model.graphs.errors import MissingBranchError
 
@@ -75,58 +76,41 @@ def get_downstream_nodes(grid: "Grid", node_id: int, inclusive: bool = False):
     )
 
 
-_THREE_WINDING_BRANCH_CONFIGS = (
-    ("node_1", "node_2", "status_1", "status_2"),
-    ("node_1", "node_3", "status_1", "status_3"),
-    ("node_2", "node_3", "status_2", "status_3"),
-)
-
-
-def _active_branches_for_path(
-    grid: "Grid", path_nodes: list[int]
-) -> tuple[BranchArray, dict[tuple[int, int], list[int]]]:
+def _get_branch(grid: "Grid", from_node: int, to_node: int) -> BranchArray:
     """Return active branch records and an index filtered to the requested path nodes."""
 
-    active = grid.branches.filter(from_status=1, to_status=1).filter(
-        from_node=path_nodes, to_node=path_nodes, mode_="AND"
+    active_branches = grid.branches.filter(from_status=1, to_status=1).filter(
+        from_node=from_node, to_node=to_node, mode_="AND"
     )
     if grid.three_winding_transformer.size:
-        three_winding_active = (
-            grid.three_winding_transformer.as_branches()
-            .filter(from_status=1, to_status=1)
-            .filter(from_node=path_nodes, to_node=path_nodes, mode_="AND")
+        three_winding_active = grid.three_winding_transformer.as_branches().filter(
+            from_status=1, to_status=1, from_node=from_node, to_node=to_node, mode_="AND"
         )
         if three_winding_active.size:
-            active = fp.concatenate(active, three_winding_active)
-
-    index: dict[tuple[int, int], list[int]] = defaultdict(list)
-    for position, (source, target) in enumerate(zip(active.from_node, active.to_node)):
-        index[(int(source), int(target))].append(position)
+            active_branches = fp.concatenate(active_branches, three_winding_active)
 
-    return active, index
+    return active_branches
 
 
 def iter_branches_in_shortest_path(
     grid: "Grid", from_node_id: int, to_node_id: int, typed: bool = False
-) -> Iterator[BranchArray]:
+) -> Iterator[BranchArray] | Iterator[BranchArray | Branch3Array]:
     """See Grid.iter_branches_in_shortest_path()."""
 
     path, _ = grid.graphs.active_graph.get_shortest_path(from_node_id, to_node_id)
-    active_branches, index = _active_branches_for_path(grid, path)
 
-    for current_node, next_node in zip(path[:-1], path[1:]):
-        positions = index.get((current_node, next_node))
-        if not positions:
+    for current_node, next_node in pairwise(path):
+        branches = _get_branch(grid, current_node, next_node)
+        if branches.size == 0:
             raise MissingBranchError(
                 f"No active branch connects nodes {current_node} -> {next_node} even though a path exists."
             )
-        branch_records = active_branches[positions]
         if typed:
-            branch_ids = branch_records.id.tolist()
+            branch_ids = branches.id.tolist()
             try:
                 typed_branches = grid.get_typed_branches(branch_ids)
             except RecordDoesNotExist:
                 typed_branches = grid.three_winding_transformer.filter(branch_ids)
             yield typed_branches
         else:
-            yield branch_records
+            yield branches

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

@@ -61,9 +61,7 @@
     get_downstream_nodes,
     get_nearest_substation_node,
     get_typed_branches,
-)
-from power_grid_model_ds._core.model.grids._search import (
-    iter_branches_in_shortest_path as _iter_branches_in_shortest_path,
+    iter_branches_in_shortest_path,
 )
 from power_grid_model_ds._core.model.grids.serialization.json import deserialize_from_json, serialize_to_json
 from power_grid_model_ds._core.model.grids.serialization.pickle import load_grid_from_pickle, save_grid_to_pickle
@@ -327,7 +325,7 @@ def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
 
     def iter_branches_in_shortest_path(
         self, from_node_id: int, to_node_id: int, typed: bool = False
-    ) -> Iterator[BranchArray]:
+    ) -> Iterator[BranchArray] | Iterator[BranchArray | Branch3Array]:
         """Returns the ordered active branches that form the shortest path between two nodes. When parallel active edges
         are in the path all these branches will be returned for the same from_node and to_node.
 
@@ -344,7 +342,7 @@ def iter_branches_in_shortest_path(
             MissingBranchError: If the graph reports an edge on the shortest path but no active branch is found.
         """
 
-        return _iter_branches_in_shortest_path(self, from_node_id, to_node_id, typed=typed)
+        return iter_branches_in_shortest_path(self, from_node_id, to_node_id, typed=typed)
 
     def get_nearest_substation_node(self, node_id: int):
         """Find the nearest substation node.

tests/unit/model/grids/test_search.py

@@ -6,7 +6,6 @@
 import pytest
 
 from power_grid_model_ds import Grid
-from power_grid_model_ds._core.model.arrays import LineArray, ThreeWindingTransformerArray, TransformerArray
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 
@@ -60,31 +59,26 @@ def test_get_branches_in_path_empty_path(self, basic_grid):
 class TestIterBranchesInShortestPath:
     def test_iter_branches_in_shortest_path_returns_branch_arrays(self, basic_grid):
         branches = list(basic_grid.iter_branches_in_shortest_path(101, 106))
-        assert len(branches) == 2
-        branch_nodes = [(branch.from_node.item(), branch.to_node.item()) for branch in branches]
-        assert branch_nodes == [(101, 102), (102, 106)]
+        assert branches == [basic_grid.branches.filter(id=201), basic_grid.branches.filter(id=301)]
 
     def test_iter_branches_in_shortest_path_three_winding_transformer(self, grid_with_3wt):
         branches = list(grid_with_3wt.iter_branches_in_shortest_path(101, 104))
         assert len(branches) == 2
-        branch_nodes = [(branch.from_node.item(), branch.to_node.item()) for branch in branches]
-        assert branch_nodes == [(101, 102), (102, 104)]
+        assert branches[0].id.item() == 301
+        assert branches[0].from_node.item() == 101
+        assert branches[0].to_node.item() == 102
+        assert branches[1] == grid_with_3wt.branches.filter(id=201)
 
     def test_iter_branches_same_node_returns_empty(self, basic_grid):
         assert [] == list(basic_grid.iter_branches_in_shortest_path(101, 101))
 
     def test_iter_branches_in_shortest_path_typed(self, basic_grid):
         branches = list(basic_grid.iter_branches_in_shortest_path(101, 106, typed=True))
-        assert len(branches) == 2
-        assert isinstance(branches[0], LineArray)
-        assert isinstance(branches[1], TransformerArray)
-        assert [branch.id.item() for branch in branches] == [201, 301]
+        assert branches == [basic_grid.line.filter(id=201), basic_grid.transformer.filter(id=301)]
 
     def test_iter_branches_in_shortest_path_three_winding_transformer_typed(self, grid_with_3wt):
         branches = list(grid_with_3wt.iter_branches_in_shortest_path(101, 104, typed=True))
-        assert len(branches) == 2
-        assert isinstance(branches[0], ThreeWindingTransformerArray)
-        assert isinstance(branches[1], LineArray)
+        assert branches == [grid_with_3wt.three_winding_transformer.filter(id=301), grid_with_3wt.line.filter(id=201)]
 
 
 def test_component_three_winding_transformer(grid_with_3wt):