Add Branch3.as_branches method

Signed-off-by: Thijs Baaijen <[email protected]>

Author: Thijss

src/power_grid_model_ds/_core/model/arrays/base/array.py

@@ -66,7 +66,7 @@ def __init__(self: Self, *args, data: NDArray | None = None, **kwargs):
             self._data = data
 
     @property
-    def data(self: Self):
+    def data(self: Self) -> NDArray:
         return self._data
 
     @classmethod

src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py

@@ -9,6 +9,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from power_grid_model_ds._core.fancypy import concatenate
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.dtypes.appliances import Source, SymGen, SymLoad
 from power_grid_model_ds._core.model.dtypes.branches import (
@@ -99,7 +100,26 @@ class TransformerArray(Transformer, BranchArray):
 
 
 class Branch3Array(IdArray, Branch3):
-    pass
+    def as_branches(self) -> BranchArray:
+        """Convert Branch3Array to BranchArray."""
+        branches_1_2 = BranchArray.empty(self.size)
+        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.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.from_node = self.node_2
+        branches_2_3.to_node = self.node_3
+        branches_2_3.from_status = self.status_2
+        branches_2_3.to_status = self.status_3
+        return concatenate(branches_1_2, branches_1_3, branches_2_3)
 
 
 class ThreeWindingTransformerArray(Branch3Array, ThreeWindingTransformer):

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

@@ -15,7 +15,6 @@
     MissingNodeError,
     NoPathBetweenNodes,
 )
-from power_grid_model_ds._core.model.utils import _get_branch3_branches
 
 if TYPE_CHECKING:
     from power_grid_model_ds._core.model.grids.base import Grid
@@ -175,8 +174,7 @@ def add_branch_array(self, branch_array: BranchArray) -> None:
     def add_branch3_array(self, branch3_array: Branch3Array) -> None:
         """Add all branch3s in the branch3 array to the graph."""
         for branch3 in branch3_array:
-            branches = _get_branch3_branches(branch3)
-            self.add_branch_array(branches)
+            self.add_branch_array(branch3.as_branches())
 
     def delete_branch_array(self, branch_array: BranchArray, raise_on_fail: bool = True) -> None:
         """Delete all branches in branch_array from the graph."""
@@ -187,8 +185,7 @@ def delete_branch_array(self, branch_array: BranchArray, raise_on_fail: bool = T
     def delete_branch3_array(self, branch3_array: Branch3Array, raise_on_fail: bool = True) -> None:
         """Delete all branch3s in the branch3 array from the graph."""
         for branch3 in branch3_array:
-            branches = _get_branch3_branches(branch3)
-            self.delete_branch_array(branches, raise_on_fail=raise_on_fail)
+            self.delete_branch_array(branch3.as_branches(), raise_on_fail=raise_on_fail)
 
     @contextmanager
     def tmp_remove_nodes(self, nodes: list[int]) -> Generator:

src/power_grid_model_ds/_core/model/utils.py

@@ -6,7 +6,6 @@
 
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.grids.base import Grid
-from power_grid_model_ds._core.model.utils import _get_branch3_branches
 from power_grid_model_ds.arrays import Branch3Array, BranchArray, NodeArray
 
 
@@ -41,15 +40,15 @@ def parse_branch3_array(branches: Branch3Array, group: Literal["transformer"]) -
     """Parse the three-winding transformer array."""
     parsed_branches = []
     columns = branches.columns
-    for branch in branches:
-        for branch_ in _get_branch3_branches(branch):
-            cyto_elements = {"data": _array_to_dict(branch_, columns)}
+    for branch3 in branches:
+        for branch1 in branch3.as_branches():
+            cyto_elements = {"data": _array_to_dict(branch1, columns)}
             cyto_elements["data"].update(
                 {
                     # IDs need to be unique, so we combine the branch ID with the from and to nodes
-                    "id": str(branch.id.item()) + f"_{branch_.from_node.item()}_{branch_.to_node.item()}",
-                    "source": str(branch_.from_node.item()),
-                    "target": str(branch_.to_node.item()),
+                    "id": str(branch3.id.item()) + f"_{branch1.from_node.item()}_{branch1.to_node.item()}",
+                    "source": str(branch1.from_node.item()),
+                    "target": str(branch1.to_node.item()),
                     "group": group,
                 }
             )

src/power_grid_model_ds/_core/visualizer/parsers.py

@@ -5,7 +5,7 @@
 import pytest
 from numpy.testing import assert_array_equal
 
-from power_grid_model_ds._core.model.arrays import BranchArray
+from power_grid_model_ds._core.model.arrays import Branch3Array, BranchArray
 
 # pylint: disable=missing-function-docstring
 
@@ -18,28 +18,66 @@ def parallel_branches():
     return branches
 
 
-def test_branch_is_active():
-    branches = BranchArray.zeros(4)
-    branches.from_status = [1, 1, 0, 0]
-    branches.to_status = [1, 0, 1, 0]
+class TestBranchArray:
+    def test_branch_is_active(self):
+        branches = BranchArray.zeros(4)
+        branches.from_status = [1, 1, 0, 0]
+        branches.to_status = [1, 0, 1, 0]
 
-    assert_array_equal(branches.is_active, [True, False, False, False])
-    assert branches[0].is_active
+        assert_array_equal(branches.is_active, [True, False, False, False])
+        assert branches[0].is_active
 
+    def test_branch_node_ids(self):
+        branches = BranchArray.zeros(2)
+        branches.from_node = [0, 1]
+        branches.to_node = [1, 2]
 
-def test_branch_node_ids():
-    branches = BranchArray.zeros(2)
-    branches.from_node = [0, 1]
-    branches.to_node = [1, 2]
+        assert_array_equal(branches.node_ids, [0, 1, 1, 2])
 
-    assert_array_equal(branches.node_ids, [0, 1, 1, 2])
+    def test_filter_non_parallel(self, branches: BranchArray):
+        filtered_branches = branches.filter_parallel(1, "eq")
+        assert_array_equal(filtered_branches.data, branches[2].data)
 
+    def test_filter_parallel(self, branches: BranchArray):
+        filtered_branches = branches.filter_parallel(1, "neq")
+        assert_array_equal(filtered_branches.data, branches[0:2].data)
 
-def test_filter_non_parallel(branches: BranchArray):
-    filtered_branches = branches.filter_parallel(1, "eq")
-    assert_array_equal(filtered_branches.data, branches[2].data)
 
+class TestBranch3Array:
+    def test_as_branches_single(self):
+        branch3 = Branch3Array(
+            node_1=[1],
+            node_2=[2],
+            node_3=[3],
+            status_1=[1],
+            status_2=[1],
+            status_3=[0],
+        )
 
-def test_filter_parallel(branches: BranchArray):
-    filtered_branches = branches.filter_parallel(1, "neq")
-    assert_array_equal(filtered_branches.data, branches[0:2].data)
+        branch_array = branch3.as_branches()
+
+        assert branch_array.size == 3
+
+        assert branch_array.from_node.tolist() == [1, 1, 2]
+        assert branch_array.to_node.tolist() == [2, 3, 3]
+        assert branch_array.from_status.tolist() == [1, 1, 1]
+        assert branch_array.to_status.tolist() == [1, 0, 0]
+
+    def test_as_branches_multiple(self):
+        branch3 = Branch3Array(
+            node_1=[1, 4],
+            node_2=[2, 5],
+            node_3=[3, 6],
+            status_1=[1, 1],
+            status_2=[1, 1],
+            status_3=[1, 0],
+        )
+
+        branch_array = branch3.as_branches()
+        branch_array.sort(order=["from_node", "to_node"])
+
+        assert branch_array.size == 6
+        assert branch_array.from_node.tolist() == [1, 1, 2, 4, 4, 5]
+        assert branch_array.to_node.tolist() == [2, 3, 3, 5, 6, 6]
+        assert branch_array.from_status.tolist() == [1, 1, 1, 1, 1, 1]
+        assert branch_array.to_status.tolist() == [1, 1, 1, 1, 0, 0]