support adding multiple nodes at once to the graph

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

Author: Thijss

src/power_grid_model_ds/_core/model/graphs/container.py

@@ -153,6 +153,8 @@ def _append(self, array: FancyArray) -> None:
         if isinstance(array, BranchArray):
             self.add_branch(array)
         if isinstance(array, Branch3Array):
-            self.add_branch3(array)
+            for record in array:
+                self.add_branch3(record)
         if isinstance(array, NodeArray):
-            self.add_node(array)
+            for record in array:
+                self.add_node(record)

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

@@ -72,6 +72,7 @@ def add_node(self, ext_node_id: int, raise_on_fail: bool = True) -> None:
 
         self._add_node(ext_node_id)
 
+
     def delete_node(self, ext_node_id: int, raise_on_fail: bool = True) -> None:
         """Remove a node from the graph.
 
@@ -93,8 +94,10 @@ def delete_node(self, ext_node_id: int, raise_on_fail: bool = True) -> None:
 
     def add_node_array(self, node_array: NodeArray, raise_on_fail: bool = True) -> None:
         """Add all nodes in the node array to the graph."""
-        for node in node_array:
-            self.add_node(ext_node_id=node.id.item(), raise_on_fail=raise_on_fail)
+        ext_node_ids = node_array.id.tolist()
+        if existing_ids := set(ext_node_ids).intersection(set(self.external_ids)):
+            raise GraphError(f"{len(existing_ids)} external node ids already exist!")
+        self._add_nodes(ext_node_ids)
 
     def delete_node_array(self, node_array: NodeArray, raise_on_fail: bool = True) -> None:
         """Delete all nodes in node_array from the graph"""
@@ -142,9 +145,14 @@ def delete_branch(self, from_ext_node_id: int, to_ext_node_id: int, raise_on_fai
 
     def add_branch_array(self, branch_array: BranchArray) -> None:
         """Add all branches in the branch array to the graph."""
-        for branch in branch_array:
-            if self._branch_is_relevant(branch):
-                self.add_branch(branch.from_node.item(), branch.to_node.item())
+        if self.active_only:
+            branch_array = branch_array[branch_array.is_active]
+            if not branch_array.size:
+                return
+
+        from_node_ids = self._externals_to_internals(branch_array.from_node.tolist())
+        to_node_ids = self._externals_to_internals(branch_array.to_node.tolist())
+        self._add_branches(from_node_ids, to_node_ids)
 
     def add_branch3_array(self, branch3_array: Branch3Array) -> None:
         """Add all branch3s in the branch3 array to the graph."""
@@ -282,11 +290,17 @@ def _has_node(self, node_id) -> bool: ...
     @abstractmethod
     def _add_node(self, ext_node_id: int) -> None: ...
 
+    @abstractmethod
+    def _add_nodes(self, ext_node_ids: list[int]) -> None: ...
+
     @abstractmethod
     def _delete_node(self, node_id: int): ...
 
     @abstractmethod
-    def _add_branch(self, from_node_id, to_node_id) -> None: ...
+    def _add_branch(self, from_node_id: int, to_node_id: int) -> None: ...
+
+    @abstractmethod
+    def _add_branches(self, from_node_ids: list[int], to_node_ids: list[int]) -> None: ...
 
     @abstractmethod
     def _delete_branch(self, from_node_id, to_node_id) -> None:

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

@@ -51,6 +51,12 @@ def _add_node(self, ext_node_id: int):
         self._external_to_internal[ext_node_id] = graph_node_id
         self._internal_to_external[graph_node_id] = ext_node_id
 
+    def _add_nodes(self, ext_node_ids: list[int]) -> None:
+        graph_node_ids = self._graph.add_nodes_from(ext_node_ids)
+        for ext_node_id, graph_node_id in zip(ext_node_ids, graph_node_ids):
+            self._external_to_internal[ext_node_id] = graph_node_id
+            self._internal_to_external[graph_node_id] = ext_node_id
+
     def _delete_node(self, node_id: int):
         self._graph.remove_node(node_id)
         external_node_id = self._internal_to_external.pop(node_id)
@@ -65,6 +71,10 @@ def _has_node(self, node_id: int) -> bool:
     def _add_branch(self, from_node_id: int, to_node_id: int):
         self._graph.add_edge(from_node_id, to_node_id, None)
 
+    def _add_branches(self, from_node_ids: list[int], to_node_ids: list[int]):
+        edge_list = [(from_node_id, to_node_id, None) for from_node_id, to_node_id in zip(from_node_ids, to_node_ids)]
+        self._graph.add_edges_from(edge_list)
+
     def _delete_branch(self, from_node_id: int, to_node_id: int) -> None:
         try:
             self._graph.remove_edge(from_node_id, to_node_id)

tests/unit/model/graphs/test_container.py

@@ -7,7 +7,6 @@
 from power_grid_model_ds._core.model.arrays import NodeArray, ThreeWindingTransformerArray
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
 from power_grid_model_ds._core.model.graphs.container import GraphContainer
-from power_grid_model_ds._core.model.graphs.errors import GraphError
 
 # pylint: disable=missing-function-docstring
 
@@ -66,23 +65,15 @@ def test_from_arrays_partially_active_three_winding(basic_grid):
     assert basic_grid.graphs.complete_graph.nr_nodes == graphs.complete_graph.nr_nodes
     assert basic_grid.graphs.complete_graph.nr_branches == 6 + 3
 
-    # Implicitly test that the correct branches are added
-    # Current implementation does not have a has_branch method.
-    basic_grid.graphs.complete_graph.delete_branch(1000, 1001)
-    basic_grid.graphs.complete_graph.delete_branch(1000, 1002)
-    basic_grid.graphs.complete_graph.delete_branch(1001, 1002)
+    basic_grid.graphs.active_graph.has_branch(1000, 1002)
+    basic_grid.graphs.active_graph.has_branch(1001, 1002)
 
     assert basic_grid.graphs.active_graph.nr_nodes == graphs.active_graph.nr_nodes
     assert basic_grid.graphs.active_graph.nr_branches == 5 + 1
 
-    # Implicitly test that the correct branches are added
-    # Current implementation does not have a has_branch method.
-    basic_grid.graphs.active_graph.delete_branch(1000, 1001)
-    with pytest.raises(GraphError):
-        basic_grid.graphs.active_graph.delete_branch(1000, 1002)
-
-    with pytest.raises(GraphError):
-        basic_grid.graphs.active_graph.delete_branch(1001, 1002)
+    basic_grid.graphs.active_graph.has_branch(1000, 1001)
+    assert not basic_grid.graphs.active_graph.has_branch(1000, 1002)
+    assert not basic_grid.graphs.active_graph.has_branch(1001, 1002)
 
 
 def test_from_arrays_invalid_arrays(basic_grid):

tests/unit/model/test_containers.py

@@ -156,15 +156,17 @@ def test_id_counter():
 
 
 def test_branches(grid):
-    node = NodeArray.zeros(10, empty_id=False)
-    grid.append(node)
-    grid.append(LineArray.zeros(10))
-    grid.append(TransformerArray.zeros(10))
-    grid.append(LinkArray.zeros(10))
-    branches = grid.branches
+    nodes = NodeArray.zeros(10)
+    grid.append(nodes)
+
+    for branch_class in (LineArray, TransformerArray, LinkArray):
+        branches = branch_class.zeros(10)
+        branches.from_node = nodes.id
+        branches.to_node = list(reversed(nodes.id.tolist()))
+        grid.append(branches)
 
     expected_ids = np.concatenate((grid.line.id, grid.transformer.id, grid.link.id))
-    assert set(expected_ids) == set(branches.id)
+    assert set(expected_ids) == set(grid.branches.id)
 
 
 def test_delete_node_without_additional_properties(basic_grid):