Feat: support adding multiple nodes/branches at once (performance increase)

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

@@ -55,18 +55,32 @@ def empty(cls, graph_model: type[BaseGraphModel] = RustworkxGraphModel) -> "Grap
             complete_graph=graph_model(active_only=False),
         )
 
-    def add_branch(self, branch: BranchArray) -> None:
+    def add_node_array(self, node_array: NodeArray) -> None:
+        """Add a node to all graphs"""
+        for field in dataclasses.fields(self):
+            graph = getattr(self, field.name)
+            graph.add_node_array(node_array=node_array, raise_on_fail=False)
+            setattr(self, field.name, graph)
+
+    def add_branch_array(self, branch_array: BranchArray) -> None:
         """Add a branch to all graphs"""
         for field in self.graph_attributes:
             graph = getattr(self, field.name)
-            graph.add_branch_array(branch_array=branch)
+            graph.add_branch_array(branch_array=branch_array)
             setattr(self, field.name, graph)
 
-    def add_branch3(self, branch: Branch3Array) -> None:
+    def add_branch3_array(self, branch3_array: Branch3Array) -> None:
         """Add a branch to all graphs"""
         for field in self.graph_attributes:
             graph = getattr(self, field.name)
-            graph.add_branch3_array(branch3_array=branch)
+            graph.add_branch3_array(branch3_array=branch3_array)
+            setattr(self, field.name, graph)
+
+    def delete_node(self, node: NodeArray) -> None:
+        """Remove a node from all graphs"""
+        for field in dataclasses.fields(self):
+            graph = getattr(self, field.name)
+            graph.delete_node_array(node_array=node)
             setattr(self, field.name, graph)
 
     def delete_branch(self, branch: BranchArray) -> None:
@@ -83,20 +97,6 @@ def delete_branch3(self, branch: Branch3Array) -> None:
             graph.delete_branch3_array(branch3_array=branch)
             setattr(self, field.name, graph)
 
-    def add_node(self, node: NodeArray) -> None:
-        """Add a node to all graphs"""
-        for field in dataclasses.fields(self):
-            graph = getattr(self, field.name)
-            graph.add_node_array(node_array=node, raise_on_fail=False)
-            setattr(self, field.name, graph)
-
-    def delete_node(self, node: NodeArray) -> None:
-        """Remove a node from all graphs"""
-        for field in dataclasses.fields(self):
-            graph = getattr(self, field.name)
-            graph.delete_node_array(node_array=node)
-            setattr(self, field.name, graph)
-
     def make_active(self, branch: BranchArray) -> None:
         """Add branch to all active_only graphs"""
 
@@ -150,9 +150,10 @@ def _validate_branches(arrays: MinimalGridArrays):
                 raise RecordDoesNotExist(f"Found invalid .to_node values in {array.__class__.__name__}")
 
     def _append(self, array: FancyArray) -> None:
+        if isinstance(array, NodeArray):
+            self.add_node_array(array)
         if isinstance(array, BranchArray):
-            self.add_branch(array)
+            self.add_branch_array(array)
         if isinstance(array, Branch3Array):
-            self.add_branch3(array)
-        if isinstance(array, NodeArray):
-            self.add_node(array)
+            for record in array:
+                self.add_branch3_array(record)

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

@@ -93,8 +93,11 @@ 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 raise_on_fail:
+            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)

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

@@ -200,9 +200,9 @@ def append(self, array: FancyArray, check_max_id: bool = True):
             check_max_id (bool, optional): Whether to check if the array id is the maximum id. Defaults to True.
         """
         self._append(array, check_max_id=check_max_id)  # noqa
-        for row in array:
-            # pylint: disable=protected-access
-            self.graphs._append(row)
+
+        # pylint: disable=protected-access
+        self.graphs._append(array)
 
     def add_branch(self, branch: BranchArray) -> None:
         """Add a branch to the grid
@@ -211,7 +211,7 @@ def add_branch(self, branch: BranchArray) -> None:
             branch (BranchArray): The branch to add
         """
         self._append(array=branch)
-        self.graphs.add_branch(branch=branch)
+        self.graphs.add_branch_array(branch_array=branch)
 
         logging.debug(f"added branch {branch.id} from {branch.from_node} to {branch.to_node}")
 
@@ -243,7 +243,7 @@ def add_node(self, node: NodeArray) -> None:
             node (NodeArray): The node to add
         """
         self._append(array=node)
-        self.graphs.add_node(node=node)
+        self.graphs.add_node_array(node_array=node)
         logging.debug(f"added rail {node.id}")
 
     def delete_node(self, node: NodeArray) -> None:

tests/performance/grid_performance_tests.py

@@ -7,7 +7,37 @@
 # pylint: disable=missing-function-docstring
 
 
-def test_get_downstream_nodes_performance():
+def perf_test_add_nodes():
+    setup_code = {
+        "grid": "from power_grid_model_ds import Grid;"
+        + "from power_grid_model_ds._core.model.arrays import NodeArray;"
+        + "grid = Grid.empty();"
+        + "nodes = NodeArray.zeros({size});"
+    }
+
+    code_to_test = ["grid.append(nodes);"]
+
+    do_performance_test(code_to_test, [10, 200, 1000], 100, setup_code)
+
+
+def perf_test_add_lines():
+    setup_code = {
+        "grid": "from power_grid_model_ds import Grid;"
+        + "from power_grid_model_ds._core.model.arrays import NodeArray, LineArray;"
+        + "grid = Grid.empty();"
+        + "nodes = NodeArray.zeros({size});"
+        + "grid.append(nodes);"
+        + "lines = LineArray.zeros({size});"
+        + "lines.from_node = nodes.id;"
+        + "lines.to_node = nodes.id;"
+    }
+
+    code_to_test = ["grid.append(lines);"]
+
+    do_performance_test(code_to_test, [10, 200, 1000], 100, setup_code)
+
+
+def perf_test_get_downstream_nodes_performance():
     setup_code = {
         "grid": "import numpy as np;"
         + "from power_grid_model_ds.enums import NodeType;"
@@ -27,4 +57,6 @@ def test_get_downstream_nodes_performance():
 
 
 if __name__ == "__main__":
-    test_get_downstream_nodes_performance()
+    perf_test_get_downstream_nodes_performance()
+    perf_test_add_nodes()
+    perf_test_add_lines()

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):