From 28f8e3e335cb8e6f56da9b0542646f59e432e167 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Wed, 11 Feb 2026 20:40:37 +0100
Subject: [PATCH 01/18] feat: iter branches

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 .../_core/model/grids/base.py                 | 44 ++++++++++++++++++-
 tests/unit/model/grids/test_search.py         | 17 +++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index ea877fd2..8a4eab98 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -5,13 +5,14 @@
 """Base grid classes"""
 
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from pathlib import Path
-from typing import Literal, Self, Type, TypeVar
+from typing import Iterator, Literal, Self, Type, TypeVar
 
 import numpy as np
 import numpy.typing as npt
 
+from power_grid_model_ds._core.fancypy import concatenate
 from power_grid_model_ds._core.model.arrays import (
     AsymCurrentSensorArray,
     AsymLineArray,
@@ -36,6 +37,7 @@
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.containers.base import FancyArrayContainer
 from power_grid_model_ds._core.model.graphs.container import GraphContainer
+from power_grid_model_ds._core.model.graphs.errors import MissingBranchError
 from power_grid_model_ds._core.model.graphs.models import RustworkxGraphModel
 from power_grid_model_ds._core.model.graphs.models.base import BaseGraphModel
 from power_grid_model_ds._core.model.grids._feeders import set_feeder_ids
@@ -122,6 +124,17 @@ def __str__(self) -> str:
         """
         return serialize_to_str(self)
 
+    def __repr__(self) -> str:
+        """Expose non-empty arrays with their field names for debugging."""
+        array_reprs: list[str] = []
+        for field in fields(self):
+            value = getattr(self, field.name)
+            if isinstance(value, FancyArray) and len(value):
+                array_reprs.append(f"{field.name}={value!r}")
+        if not array_reprs:
+            return f"{self.__class__.__name__}()"
+        return f"{self.__class__.__name__}({', '.join(array_reprs)})"
+
     @classmethod
     def empty(cls: Type[G], graph_model: type[BaseGraphModel] = RustworkxGraphModel) -> G:
         """Create an empty grid
@@ -311,6 +324,19 @@ def reverse_branches(self, branches: BranchArray):
         """Reverse the direction of the branches."""
         return reverse_branches(self, branches)
 
+    def _active_branches(self) -> BranchArray:
+        """Collection of active branch records including converted three-winding transformer edges."""
+
+        active = self.branches.filter(from_status=1, to_status=1)
+        if not self.three_winding_transformer.size:
+            return active
+
+        three_winding_active = self.three_winding_transformer.as_branches().filter(from_status=1, to_status=1)
+        if not three_winding_active.size:
+            return active
+
+        return concatenate(active, three_winding_active)
+
     def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
         """Returns all branches within a path of nodes
 
@@ -322,6 +348,20 @@ def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
         """
         return self.branches.filter(from_node=nodes_in_path, to_node=nodes_in_path, from_status=1, to_status=1)
 
+    def iter_branches_in_shortest_path(self, from_node_id: int, to_node_id: int) -> Iterator[BranchArray]:
+        """Yield the active branches that connect two nodes via the shortest path."""
+
+        path, _ = self.graphs.active_graph.get_shortest_path(from_node_id, to_node_id)
+        active_branches = self._active_branches()
+
+        for current_node, next_node in zip(path[:-1], path[1:]):
+            branch = active_branches.filter(from_node=[current_node], to_node=[next_node])
+            if not len(branch):
+                raise MissingBranchError(
+                    f"No active branch connects nodes {current_node} -> {next_node} even though a path exists."
+                )
+            yield branch
+
     def get_nearest_substation_node(self, node_id: int):
         """Find the nearest substation node.
 
diff --git a/tests/unit/model/grids/test_search.py b/tests/unit/model/grids/test_search.py
index f42891b6..0cb3678a 100644
--- a/tests/unit/model/grids/test_search.py
+++ b/tests/unit/model/grids/test_search.py
@@ -56,6 +56,23 @@ def test_get_branches_in_path_empty_path(self, basic_grid):
         assert 0 == branches.size
 
 
+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)]
+
+    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)]
+
+    def test_iter_branches_same_node_returns_empty(self, basic_grid):
+        assert [] == list(basic_grid.iter_branches_in_shortest_path(101, 101))
+
+
 def test_component_three_winding_transformer(grid_with_3wt):
     substation_nodes = grid_with_3wt.node.filter(node_type=NodeType.SUBSTATION_NODE.value).id
     with grid_with_3wt.graphs.active_graph.tmp_remove_nodes(substation_nodes):

From 102c2809e358eafa05b95b1ab520a04699f47da8 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Wed, 11 Feb 2026 20:56:31 +0100
Subject: [PATCH 02/18] feat: add grid repr method

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/base.py | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index 8a4eab98..cd163d44 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -5,7 +5,7 @@
 """Base grid classes"""
 
 import warnings
-from dataclasses import dataclass, fields
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterator, Literal, Self, Type, TypeVar
 
@@ -124,17 +124,6 @@ def __str__(self) -> str:
         """
         return serialize_to_str(self)
 
-    def __repr__(self) -> str:
-        """Expose non-empty arrays with their field names for debugging."""
-        array_reprs: list[str] = []
-        for field in fields(self):
-            value = getattr(self, field.name)
-            if isinstance(value, FancyArray) and len(value):
-                array_reprs.append(f"{field.name}={value!r}")
-        if not array_reprs:
-            return f"{self.__class__.__name__}()"
-        return f"{self.__class__.__name__}({', '.join(array_reprs)})"
-
     @classmethod
     def empty(cls: Type[G], graph_model: type[BaseGraphModel] = RustworkxGraphModel) -> G:
         """Create an empty grid

From ff85f14dd1b6ecea660d0d38dc1bbe736ef8223b Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Wed, 11 Feb 2026 21:05:43 +0100
Subject: [PATCH 03/18] feat: add grid repr method

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 .../_core/model/grids/_search.py              | 35 ++++++++++++++++++-
 .../_core/model/grids/base.py                 | 33 ++++-------------
 2 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index 949bf71b..d1ae5c15 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -3,7 +3,8 @@
 # SPDX-License-Identifier: MPL-2.0
 
 import dataclasses
-from typing import TYPE_CHECKING
+from collections import defaultdict
+from typing import TYPE_CHECKING, Iterator
 
 import numpy as np
 import numpy.typing as npt
@@ -12,6 +13,7 @@
 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.enums.nodes import NodeType
+from power_grid_model_ds._core.model.graphs.errors import MissingBranchError
 
 if TYPE_CHECKING:
     from power_grid_model_ds._core.model.grids.base import Grid
@@ -71,3 +73,34 @@ def get_downstream_nodes(grid: "Grid", node_id: int, inclusive: bool = False):
     return grid.graphs.active_graph.get_downstream_nodes(
         node_id=node_id, start_node_ids=list(substation_nodes.id), inclusive=inclusive
     )
+
+
+def _active_branches(grid: "Grid") -> tuple[BranchArray, dict[tuple[int, int], list[int]]]:
+    """Return active branch records plus an index keyed on their node pairs."""
+
+    active = grid.branches.filter(from_status=1, to_status=1)
+    if grid.three_winding_transformer.size:
+        three_winding_active = grid.three_winding_transformer.as_branches().filter(from_status=1, to_status=1)
+        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)
+
+    return active, index
+
+
+def iter_branches_in_shortest_path(grid: "Grid", from_node_id: int, to_node_id: int) -> Iterator[BranchArray]:
+    """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(grid)
+
+    for current_node, next_node in zip(path[:-1], path[1:]):
+        positions = index.get((current_node, next_node))
+        if not positions:
+            raise MissingBranchError(
+                f"No active branch connects nodes {current_node} -> {next_node} even though a path exists."
+            )
+        yield active_branches[positions]
diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index cd163d44..d30129dc 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -12,7 +12,6 @@
 import numpy as np
 import numpy.typing as npt
 
-from power_grid_model_ds._core.fancypy import concatenate
 from power_grid_model_ds._core.model.arrays import (
     AsymCurrentSensorArray,
     AsymLineArray,
@@ -37,7 +36,6 @@
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.containers.base import FancyArrayContainer
 from power_grid_model_ds._core.model.graphs.container import GraphContainer
-from power_grid_model_ds._core.model.graphs.errors import MissingBranchError
 from power_grid_model_ds._core.model.graphs.models import RustworkxGraphModel
 from power_grid_model_ds._core.model.graphs.models.base import BaseGraphModel
 from power_grid_model_ds._core.model.grids._feeders import set_feeder_ids
@@ -64,6 +62,9 @@
     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,
+)
 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
 from power_grid_model_ds._core.model.grids.serialization.string import (
@@ -313,19 +314,6 @@ def reverse_branches(self, branches: BranchArray):
         """Reverse the direction of the branches."""
         return reverse_branches(self, branches)
 
-    def _active_branches(self) -> BranchArray:
-        """Collection of active branch records including converted three-winding transformer edges."""
-
-        active = self.branches.filter(from_status=1, to_status=1)
-        if not self.three_winding_transformer.size:
-            return active
-
-        three_winding_active = self.three_winding_transformer.as_branches().filter(from_status=1, to_status=1)
-        if not three_winding_active.size:
-            return active
-
-        return concatenate(active, three_winding_active)
-
     def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
         """Returns all branches within a path of nodes
 
@@ -338,18 +326,9 @@ def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
         return self.branches.filter(from_node=nodes_in_path, to_node=nodes_in_path, from_status=1, to_status=1)
 
     def iter_branches_in_shortest_path(self, from_node_id: int, to_node_id: int) -> Iterator[BranchArray]:
-        """Yield the active branches that connect two nodes via the shortest path."""
-
-        path, _ = self.graphs.active_graph.get_shortest_path(from_node_id, to_node_id)
-        active_branches = self._active_branches()
-
-        for current_node, next_node in zip(path[:-1], path[1:]):
-            branch = active_branches.filter(from_node=[current_node], to_node=[next_node])
-            if not len(branch):
-                raise MissingBranchError(
-                    f"No active branch connects nodes {current_node} -> {next_node} even though a path exists."
-                )
-            yield branch
+        """Yield all active branches on the shortest path connecting ``from_node_id`` to ``to_node_id``."""
+
+        return _iter_branches_in_shortest_path(self, from_node_id, to_node_id)
 
     def get_nearest_substation_node(self, node_id: int):
         """Find the nearest substation node.

From 5baeaf71ffa70f85f03f550178e8f4f9042e7b08 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Thu, 12 Feb 2026 20:01:27 +0100
Subject: [PATCH 04/18] chore: docs

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/base.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index d30129dc..dc118c68 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -326,7 +326,18 @@ def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
         return self.branches.filter(from_node=nodes_in_path, to_node=nodes_in_path, from_status=1, to_status=1)
 
     def iter_branches_in_shortest_path(self, from_node_id: int, to_node_id: int) -> Iterator[BranchArray]:
-        """Yield all active branches on the shortest path connecting ``from_node_id`` to ``to_node_id``."""
+        """Returns the ordered active branches that form the shortest path between two nodes.
+
+        Args:
+            from_node_id (int): External id of the path start node.
+            to_node_id (int): External id of the path end node.
+
+        Yields:
+            BranchArray: Single-row branch arrays for each active branch on the path.
+
+        Raises:
+            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)
 

From 1761bb8dd2b894d0c98834ae9fe1a8a18509e5de Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Thu, 12 Feb 2026 20:07:14 +0100
Subject: [PATCH 05/18] chore: docs

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 docs/examples/model/graph_examples.ipynb | 30 ++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/docs/examples/model/graph_examples.ipynb b/docs/examples/model/graph_examples.ipynb
index fc92f7ba..33cb4587 100644
--- a/docs/examples/model/graph_examples.ipynb
+++ b/docs/examples/model/graph_examples.ipynb
@@ -15,7 +15,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -37,7 +37,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -60,7 +60,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -105,6 +105,28 @@
     "print(f\"Shortest path: {path}, Length: {length}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Branches on the Shortest Path\n",
+    "\n",
+    "`Grid.iter_branches_in_shortest_path` walks the same nodes returned by `get_shortest_path` but exposes the actual `BranchArray` records for each edge. Iterate the result to inspect branch IDs, statuses, or any other metadata without recomputing the path."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from power_grid_model_ds import Grid\n",
+    "\n",
+    "grid = Grid.from_txt(\"S1 101\", \"101 102\", \"102 103\")\n",
+    "for branch in grid.iter_branches_in_shortest_path(101, 103):\n",
+    "    print(f\"Branch {branch.id.item()} runs {branch.from_node.item()} → {branch.to_node.item()}\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -180,7 +202,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": ".venv (3.12.6)",
    "language": "python",
    "name": "python3"
   },

From a4a1c0ec6a3d6f0873895e1d04c7a644159c7561 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Thu, 12 Feb 2026 20:15:37 +0100
Subject: [PATCH 06/18] feat: typed

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/_search.py | 11 +++++++++--
 src/power_grid_model_ds/_core/model/grids/base.py    |  8 ++++++--
 tests/unit/model/grids/test_search.py                |  8 ++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index d1ae5c15..de19f593 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -91,7 +91,9 @@ def _active_branches(grid: "Grid") -> tuple[BranchArray, dict[tuple[int, int], l
     return active, index
 
 
-def iter_branches_in_shortest_path(grid: "Grid", from_node_id: int, to_node_id: int) -> Iterator[BranchArray]:
+def iter_branches_in_shortest_path(
+    grid: "Grid", from_node_id: int, to_node_id: int, typed: bool = False
+) -> Iterator[BranchArray]:
     """See Grid.iter_branches_in_shortest_path()."""
 
     path, _ = grid.graphs.active_graph.get_shortest_path(from_node_id, to_node_id)
@@ -103,4 +105,9 @@ def iter_branches_in_shortest_path(grid: "Grid", from_node_id: int, to_node_id:
             raise MissingBranchError(
                 f"No active branch connects nodes {current_node} -> {next_node} even though a path exists."
             )
-        yield active_branches[positions]
+        branch_records = active_branches[positions]
+        if typed:
+            branch_ids = branch_records.id.tolist()
+            yield grid.get_typed_branches(branch_ids)
+        else:
+            yield branch_records
diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index dc118c68..e6df6656 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -325,12 +325,16 @@ def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
         """
         return self.branches.filter(from_node=nodes_in_path, to_node=nodes_in_path, from_status=1, to_status=1)
 
-    def iter_branches_in_shortest_path(self, from_node_id: int, to_node_id: int) -> Iterator[BranchArray]:
+    def iter_branches_in_shortest_path(
+        self, from_node_id: int, to_node_id: int, typed: bool = False
+    ) -> Iterator[BranchArray]:
         """Returns the ordered active branches that form the shortest path between two nodes.
 
         Args:
             from_node_id (int): External id of the path start node.
             to_node_id (int): External id of the path end node.
+            typed (bool): If True, each yielded branch is converted to its typed array via
+                ``get_typed_branches``.
 
         Yields:
             BranchArray: Single-row branch arrays for each active branch on the path.
@@ -339,7 +343,7 @@ def iter_branches_in_shortest_path(self, from_node_id: int, to_node_id: int) ->
             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)
+        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.
diff --git a/tests/unit/model/grids/test_search.py b/tests/unit/model/grids/test_search.py
index 0cb3678a..a0b3aba4 100644
--- a/tests/unit/model/grids/test_search.py
+++ b/tests/unit/model/grids/test_search.py
@@ -6,6 +6,7 @@
 import pytest
 
 from power_grid_model_ds import Grid
+from power_grid_model_ds._core.model.arrays import LineArray, TransformerArray
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 
@@ -72,6 +73,13 @@ def test_iter_branches_in_shortest_path_three_winding_transformer(self, grid_wit
     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]
+
 
 def test_component_three_winding_transformer(grid_with_3wt):
     substation_nodes = grid_with_3wt.node.filter(node_type=NodeType.SUBSTATION_NODE.value).id

From 5cdbb674b6966d6d3f18384d7fe6d5d287d76473 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Thu, 12 Feb 2026 20:57:09 +0100
Subject: [PATCH 07/18] feat: docs

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index e6df6656..95f1d51e 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -337,7 +337,7 @@ def iter_branches_in_shortest_path(
                 ``get_typed_branches``.
 
         Yields:
-            BranchArray: Single-row branch arrays for each active branch on the path.
+            BranchArray: branch arrays for each active branch on the path.
 
         Raises:
             MissingBranchError: If the graph reports an edge on the shortest path but no active branch is found.

From d5e2c8f02c2897a20a7ed99b25ceeda65ffd3157 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Sun, 15 Feb 2026 14:49:34 +0100
Subject: [PATCH 08/18] test: add test for typed threewinding

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/_search.py | 1 +
 tests/unit/model/grids/test_search.py                | 8 +++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index de19f593..92f94c45 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -108,6 +108,7 @@ def iter_branches_in_shortest_path(
         branch_records = active_branches[positions]
         if typed:
             branch_ids = branch_records.id.tolist()
+            # TODO: This does now not work for three winding transformers
             yield grid.get_typed_branches(branch_ids)
         else:
             yield branch_records
diff --git a/tests/unit/model/grids/test_search.py b/tests/unit/model/grids/test_search.py
index a0b3aba4..d0575be2 100644
--- a/tests/unit/model/grids/test_search.py
+++ b/tests/unit/model/grids/test_search.py
@@ -6,7 +6,7 @@
 import pytest
 
 from power_grid_model_ds import Grid
-from power_grid_model_ds._core.model.arrays import LineArray, TransformerArray
+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
 
@@ -80,6 +80,12 @@ def test_iter_branches_in_shortest_path_typed(self, basic_grid):
         assert isinstance(branches[1], TransformerArray)
         assert [branch.id.item() for branch in branches] == [201, 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)
+
 
 def test_component_three_winding_transformer(grid_with_3wt):
     substation_nodes = grid_with_3wt.node.filter(node_type=NodeType.SUBSTATION_NODE.value).id

From c98798abcdc2d9b352307747af2b3d000c5a2383 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Sun, 15 Feb 2026 15:50:50 +0100
Subject: [PATCH 09/18] docs: setup feeder ids notebook

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 .../_core/model/grids/_search.py              | 38 +++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index 92f94c45..e9935fad 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -10,7 +10,7 @@
 import numpy.typing as npt
 
 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 import BranchArray, ThreeWindingTransformerArray
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 from power_grid_model_ds._core.model.graphs.errors import MissingBranchError
@@ -75,6 +75,35 @@ 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 _lookup_three_winding_branch(grid: "Grid", node_a: int, node_b: int) -> ThreeWindingTransformerArray:
+    """Return the first active transformer that connects the node pair or raise if none exist."""
+
+    three_winding_array = grid.three_winding_transformer
+    error_message = f"No active three-winding transformer connects nodes {node_a} -> {node_b}."
+    if not three_winding_array.size:
+        raise MissingBranchError(error_message)
+
+    for node_col_a, node_col_b, status_col_a, status_col_b in _THREE_WINDING_BRANCH_CONFIGS:
+        transformer = three_winding_array.filter(
+            **{
+                node_col_a: [node_a, node_b],
+                node_col_b: [node_a, node_b],
+                status_col_a: 1,
+                status_col_b: 1,
+            }
+        )
+        if transformer.size:
+            return transformer
+    raise MissingBranchError(error_message)
+
+
 def _active_branches(grid: "Grid") -> tuple[BranchArray, dict[tuple[int, int], list[int]]]:
     """Return active branch records plus an index keyed on their node pairs."""
 
@@ -108,7 +137,10 @@ def iter_branches_in_shortest_path(
         branch_records = active_branches[positions]
         if typed:
             branch_ids = branch_records.id.tolist()
-            # TODO: This does now not work for three winding transformers
-            yield grid.get_typed_branches(branch_ids)
+            try:
+                typed_branches = grid.get_typed_branches(branch_ids)
+            except RecordDoesNotExist:
+                typed_branches = _lookup_three_winding_branch(grid, current_node, next_node)
+            yield typed_branches
         else:
             yield branch_records

From 42ec97a30c0de841b7e59dcea7a2b655523c9a07 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Sun, 15 Feb 2026 15:59:39 +0100
Subject: [PATCH 10/18] feat: active branches filtering

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 .../_core/model/grids/_search.py               | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index e9935fad..7bfd54fe 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -104,12 +104,20 @@ def _lookup_three_winding_branch(grid: "Grid", node_a: int, node_b: int) -> Thre
     raise MissingBranchError(error_message)
 
 
-def _active_branches(grid: "Grid") -> tuple[BranchArray, dict[tuple[int, int], list[int]]]:
-    """Return active branch records plus an index keyed on their node pairs."""
+def _active_branches_for_path(
+    grid: "Grid", path_nodes: list[int]
+) -> tuple[BranchArray, dict[tuple[int, int], list[int]]]:
+    """Return active branch records and an index filtered to the requested path nodes."""
 
-    active = grid.branches.filter(from_status=1, to_status=1)
+    active = grid.branches.filter(from_status=1, to_status=1).filter(
+        from_node=path_nodes, to_node=path_nodes, mode_="OR"
+    )
     if grid.three_winding_transformer.size:
-        three_winding_active = grid.three_winding_transformer.as_branches().filter(from_status=1, to_status=1)
+        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_="OR")
+        )
         if three_winding_active.size:
             active = fp.concatenate(active, three_winding_active)
 
@@ -126,7 +134,7 @@ def iter_branches_in_shortest_path(
     """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(grid)
+    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))

From 8baba8d75f98d43dd79cc0007dd1f3406449a6ad Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Sun, 15 Feb 2026 16:07:14 +0100
Subject: [PATCH 11/18] chore: mypy

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/_search.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index 7bfd54fe..3b5e636c 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -97,7 +97,7 @@ def _lookup_three_winding_branch(grid: "Grid", node_a: int, node_b: int) -> Thre
                 node_col_b: [node_a, node_b],
                 status_col_a: 1,
                 status_col_b: 1,
-            }
+            }  # type: ignore[arg-type]
         )
         if transformer.size:
             return transformer

From 0e681c6f67bb290c82bdf8383a929765ba038db3 Mon Sep 17 00:00:00 2001
From: Jaap Schouten <58551444+jaapschoutenalliander@users.noreply.github.com>
Date: Sun, 15 Feb 2026 16:14:35 +0100
Subject: [PATCH 12/18] chore: copilot review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Jaap Schouten <58551444+jaapschoutenalliander@users.noreply.github.com>
---
 src/power_grid_model_ds/_core/model/grids/_search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index 3b5e636c..d4fe1b03 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -110,13 +110,13 @@ def _active_branches_for_path(
     """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_="OR"
+        from_node=path_nodes, to_node=path_nodes, 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_="OR")
+            .filter(from_node=path_nodes, to_node=path_nodes, mode_="AND")
         )
         if three_winding_active.size:
             active = fp.concatenate(active, three_winding_active)

From 441e2c5f4e1abe043771617ce4d4d213ce2c0f97 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Sun, 15 Feb 2026 16:50:58 +0100
Subject: [PATCH 13/18] docs: clarify docstring

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/base.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index 95f1d51e..ac1ea4be 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -328,7 +328,8 @@ 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]:
-        """Returns the ordered active branches that form the shortest path between two nodes.
+        """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.
 
         Args:
             from_node_id (int): External id of the path start node.

From 5c3734bfb351de4274cd2ea74b5d1f76ee503691 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Mon, 16 Feb 2026 08:00:28 +0100
Subject: [PATCH 14/18] feat: simplify three winding lookup

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 .../_core/model/arrays/pgm_arrays.py          |  3 +++
 .../_core/model/grids/_search.py              | 26 ++-----------------
 2 files changed, 5 insertions(+), 24 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py b/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py
index 0e4b9bcf..203d18c8 100644
--- a/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py
+++ b/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py
@@ -120,18 +120,21 @@ class Branch3Array(IdArray, Branch3):
     def as_branches(self) -> BranchArray:
         """Convert Branch3Array to BranchArray."""
         branches_1_2 = BranchArray.empty(self.size)
+        branches_1_2.id = self.id
         branches_1_2.from_node = self.node_1
         branches_1_2.to_node = self.node_2
         branches_1_2.from_status = self.status_1
         branches_1_2.to_status = self.status_2
 
         branches_1_3 = BranchArray.empty(self.size)
+        branches_1_3.id = self.id
         branches_1_3.from_node = self.node_1
         branches_1_3.to_node = self.node_3
         branches_1_3.from_status = self.status_1
         branches_1_3.to_status = self.status_3
 
         branches_2_3 = BranchArray.empty(self.size)
+        branches_2_3.id = self.id
         branches_2_3.from_node = self.node_2
         branches_2_3.to_node = self.node_3
         branches_2_3.from_status = self.status_2
diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index d4fe1b03..2b08d051 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -10,7 +10,7 @@
 import numpy.typing as npt
 
 from power_grid_model_ds._core import fancypy as fp
-from power_grid_model_ds._core.model.arrays import BranchArray, ThreeWindingTransformerArray
+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.enums.nodes import NodeType
 from power_grid_model_ds._core.model.graphs.errors import MissingBranchError
@@ -82,28 +82,6 @@ def get_downstream_nodes(grid: "Grid", node_id: int, inclusive: bool = False):
 )
 
 
-def _lookup_three_winding_branch(grid: "Grid", node_a: int, node_b: int) -> ThreeWindingTransformerArray:
-    """Return the first active transformer that connects the node pair or raise if none exist."""
-
-    three_winding_array = grid.three_winding_transformer
-    error_message = f"No active three-winding transformer connects nodes {node_a} -> {node_b}."
-    if not three_winding_array.size:
-        raise MissingBranchError(error_message)
-
-    for node_col_a, node_col_b, status_col_a, status_col_b in _THREE_WINDING_BRANCH_CONFIGS:
-        transformer = three_winding_array.filter(
-            **{
-                node_col_a: [node_a, node_b],
-                node_col_b: [node_a, node_b],
-                status_col_a: 1,
-                status_col_b: 1,
-            }  # type: ignore[arg-type]
-        )
-        if transformer.size:
-            return transformer
-    raise MissingBranchError(error_message)
-
-
 def _active_branches_for_path(
     grid: "Grid", path_nodes: list[int]
 ) -> tuple[BranchArray, dict[tuple[int, int], list[int]]]:
@@ -148,7 +126,7 @@ def iter_branches_in_shortest_path(
             try:
                 typed_branches = grid.get_typed_branches(branch_ids)
             except RecordDoesNotExist:
-                typed_branches = _lookup_three_winding_branch(grid, current_node, next_node)
+                typed_branches = grid.three_winding_transformer.filter(branch_ids)
             yield typed_branches
         else:
             yield branch_records

From 6e15083f8de16ee970fe30263c8c7bf4dc3fea0f Mon Sep 17 00:00:00 2001
From: Vincent Koppen <vincent.koppen@alliander.com>
Date: Fri, 27 Feb 2026 15:44:09 +0100
Subject: [PATCH 15/18] simplify function

Signed-off-by: Vincent Koppen <vincent.koppen@alliander.com>
---
 .../_core/model/grids/_search.py              | 46 ++++++-------------
 .../_core/model/grids/base.py                 |  8 ++--
 tests/unit/model/grids/test_search.py         | 20 +++-----
 3 files changed, 25 insertions(+), 49 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index 2b08d051..d585b566 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/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
diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index ac1ea4be..a3c61c6c 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/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.
diff --git a/tests/unit/model/grids/test_search.py b/tests/unit/model/grids/test_search.py
index d0575be2..01458ed0 100644
--- a/tests/unit/model/grids/test_search.py
+++ b/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):

From 06ab70f83c156b1d6080f1072e413d5667bf9644 Mon Sep 17 00:00:00 2001
From: Vincent Koppen <vincent.koppen@alliander.com>
Date: Fri, 27 Feb 2026 16:01:34 +0100
Subject: [PATCH 16/18] naming

Signed-off-by: Vincent Koppen <vincent.koppen@alliander.com>
---
 src/power_grid_model_ds/_core/model/grids/_search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index c378e9b2..2cb73f74 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -75,7 +75,7 @@ def get_downstream_nodes(grid: "Grid", node_id: int, inclusive: bool = False):
     )
 
 
-def _get_branch(grid: "Grid", from_node: int, to_node: int) -> BranchArray:
+def _get_branches(grid: "Grid", from_node: int, to_node: int) -> BranchArray:
     """Return active branch records and an index filtered to the requested path nodes."""
 
     active_branches = grid.branches.filter(from_status=1, to_status=1).filter(
@@ -99,7 +99,7 @@ def iter_branches_in_shortest_path(
     path, _ = grid.graphs.active_graph.get_shortest_path(from_node_id, to_node_id)
 
     for current_node, next_node in pairwise(path):
-        branches = _get_branch(grid, current_node, next_node)
+        branches = _get_branches(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."

From e78ccbd15ca9a1f3e263b3a04ee24d1c55b5466c Mon Sep 17 00:00:00 2001
From: Jaap Schouten <jaap.schouten@alliander.com>
Date: Fri, 10 Apr 2026 15:31:07 +0200
Subject: [PATCH 17/18] feat: only typed

Signed-off-by: Jaap Schouten <jaap.schouten@alliander.com>
---
 .../_core/model/grids/_search.py              | 19 ++++++++-----------
 .../_core/model/grids/base.py                 |  8 +++-----
 tests/unit/model/grids/test_search.py         | 18 +++---------------
 3 files changed, 14 insertions(+), 31 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index 2cb73f74..f41dbe85 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -92,8 +92,8 @@ def _get_branches(grid: "Grid", from_node: int, to_node: int) -> BranchArray:
 
 
 def iter_branches_in_shortest_path(
-    grid: "Grid", from_node_id: int, to_node_id: int, typed: bool = False
-) -> Iterator[BranchArray] | Iterator[BranchArray | Branch3Array]:
+    grid: "Grid", from_node_id: int, to_node_id: int
+) -> Iterator[BranchArray | Branch3Array]:
     """See Grid.iter_branches_in_shortest_path()."""
 
     path, _ = grid.graphs.active_graph.get_shortest_path(from_node_id, to_node_id)
@@ -104,12 +104,9 @@ def iter_branches_in_shortest_path(
             raise MissingBranchError(
                 f"No active branch connects nodes {current_node} -> {next_node} even though a path exists."
             )
-        if typed:
-            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 branches
+        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
diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index da1063e4..e0efc07a 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -347,16 +347,14 @@ def get_branches_in_path(self, nodes_in_path: list[int]) -> BranchArray:
         return self.branches.filter(from_node=nodes_in_path, to_node=nodes_in_path, from_status=1, to_status=1)
 
     def iter_branches_in_shortest_path(
-        self, from_node_id: int, to_node_id: int, typed: bool = False
-    ) -> Iterator[BranchArray] | Iterator[BranchArray | Branch3Array]:
+        self, from_node_id: int, to_node_id: int
+    ) -> 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.
 
         Args:
             from_node_id (int): External id of the path start node.
             to_node_id (int): External id of the path end node.
-            typed (bool): If True, each yielded branch is converted to its typed array via
-                ``get_typed_branches``.
 
         Yields:
             BranchArray: branch arrays for each active branch on the path.
@@ -365,7 +363,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)
 
     def get_nearest_substation_node(self, node_id: int):
         """Find the nearest substation node.
diff --git a/tests/unit/model/grids/test_search.py b/tests/unit/model/grids/test_search.py
index c3e595a0..d583318a 100644
--- a/tests/unit/model/grids/test_search.py
+++ b/tests/unit/model/grids/test_search.py
@@ -72,27 +72,15 @@ 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):
+    def test_iter_branches_in_shortest_path(self, basic_grid):
         branches = list(basic_grid.iter_branches_in_shortest_path(101, 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
-        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)
+        assert branches == [basic_grid.line.filter(id=201), basic_grid.transformer.filter(id=301)]
 
     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 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))
+        branches = list(grid_with_3wt.iter_branches_in_shortest_path(101, 104))
         assert branches == [grid_with_3wt.three_winding_transformer.filter(id=301), grid_with_3wt.line.filter(id=201)]
 
 

From 3d894fd0c83635d6223aae5fec9a124123da12b6 Mon Sep 17 00:00:00 2001
From: Jaap Schouten <jaap.schouten@alliander.com>
Date: Fri, 10 Apr 2026 17:10:49 +0200
Subject: [PATCH 18/18] chore: order

Signed-off-by: Jaap Schouten <jaap.schouten@alliander.com>
---
 .../_core/model/grids/_search.py              | 32 +++++++++----------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/grids/_search.py b/src/power_grid_model_ds/_core/model/grids/_search.py
index 22dede42..da49c3fd 100644
--- a/src/power_grid_model_ds/_core/model/grids/_search.py
+++ b/src/power_grid_model_ds/_core/model/grids/_search.py
@@ -78,22 +78,6 @@ def get_downstream_nodes(grid: "Grid", node_id: int, inclusive: bool = False):
     )
 
 
-def _get_branches(grid: "Grid", from_node: int, to_node: int) -> BranchArray:
-    """Return active branch records and an index filtered to the requested path nodes."""
-
-    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, from_node=from_node, to_node=to_node, mode_="AND"
-        )
-        if three_winding_active.size:
-            active_branches = fp.concatenate(active_branches, three_winding_active)
-
-    return active_branches
-
-
 def iter_branches_in_shortest_path(
     grid: "Grid", from_node_id: int, to_node_id: int
 ) -> Iterator[BranchArray | Branch3Array]:
@@ -115,6 +99,22 @@ def iter_branches_in_shortest_path(
         yield typed_branches
 
 
+def _get_branches(grid: "Grid", from_node: int, to_node: int) -> BranchArray:
+    """Return active branch records and an index filtered to the requested path nodes."""
+
+    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, from_node=from_node, to_node=to_node, mode_="AND"
+        )
+        if three_winding_active.size:
+            active_branches = fp.concatenate(active_branches, three_winding_active)
+
+    return active_branches
+
+
 def find_differences_between_grids(
     grid1: "Grid", grid2: "Grid", print_diff: bool = False
 ) -> dict[str, dict[str, object]]: