feat: support merge of two grids

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

@@ -1,10 +1,24 @@
 # SPDX-FileCopyrightText: Contributors to the Power Grid Model project <[email protected]>
 #
 # SPDX-License-Identifier: MPL-2.0
+import copy
 import dataclasses
 import logging
 from typing import TYPE_CHECKING, Type, TypeVar
 
+from power_grid_model_ds._core.model.arrays import (
+    AsymVoltageSensorArray,
+    Branch3Array,
+    BranchArray,
+    IdArray,
+    NodeArray,
+    SourceArray,
+    SymGenArray,
+    SymLoadArray,
+    SymPowerSensorArray,
+    SymVoltageSensorArray,
+    TransformerTapRegulatorArray,
+)
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.graphs.container import GraphContainer
 
@@ -14,8 +28,8 @@
 if TYPE_CHECKING:
     from .base import Grid
 
-G = TypeVar("G", bound="Grid")
 
+G = TypeVar("G", bound="Grid")
 
 logger = logging.getLogger(__name__)
 
@@ -42,3 +56,54 @@ def create_empty_grid(grid_class: Type[G], graph_model: type[BaseGraphModel] = R
     empty_fields = grid_class._get_empty_fields()  # noqa # pylint: disable=protected-access
     empty_fields["graphs"] = GraphContainer.empty(graph_model=graph_model)
     return grid_class(**empty_fields)
+
+
+def merge_grids(grid: G, other_grid: G, mode: str) -> G:
+    """See Grid.merge()"""
+
+    match mode:
+        case "recalculate_ids":
+            other_grid = copy.deepcopy(other_grid)
+            offset = grid.id_counter  # Possible improvement: grid.id_counter - other_grid.min_id() + 1
+            _increment_grid_ids_by_offset(other_grid, offset)
+        case "keep_ids":
+            pass
+        case _:
+            raise NotImplementedError(f"Merge mode {mode} is not implemented")
+
+    # Append all arrays from the first grid to the second
+    for array in other_grid.all_arrays():
+        grid.append(array, check_max_id=False)
+
+    return grid
+
+
+def _increment_grid_ids_by_offset(grid: G, offset: int) -> None:
+    for array in grid.all_arrays():
+        if isinstance(array, IdArray):
+            _update_id_column(array, "id", offset)
+
+        columns = []
+        match array:
+            case SymPowerSensorArray() | SymVoltageSensorArray() | AsymVoltageSensorArray():
+                columns = []
+            case NodeArray():
+                columns = ["feeder_node_id", "feeder_branch_id"]
+            case TransformerTapRegulatorArray():
+                columns = ["regulated_object"]
+            case BranchArray():
+                columns = ["from_node", "to_node", "feeder_node_id", "feeder_branch_id"]
+            case Branch3Array():
+                columns = ["node_1", "node_2", "node_3"]
+            case SymGenArray() | SymLoadArray() | SourceArray():
+                columns = ["node"]
+            case _:
+                raise NotImplementedError(f"The array of type {type(array)} is not implemented for appending")
+
+        for column in columns:
+            _update_id_column(array, column, offset)
+
+
+def _update_id_column(array: IdArray, column: str, offset: int) -> None:
+    mask = array.is_empty(column)
+    array[column][~mask] += offset

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

@@ -36,6 +36,7 @@
 from power_grid_model_ds._core.model.grids._helpers import (
     create_empty_grid,
     create_grid_from_extended_grid,
+    merge_grids,
 )
 from power_grid_model_ds._core.model.grids._modify import (
     add_array_to_grid,
@@ -346,3 +347,17 @@ def cache(self, cache_dir: Path, cache_name: str, compress: bool = True):
             compress (bool, optional): Whether to compress the cache. Defaults to True.
         """
         return save_grid_to_pickle(self, cache_dir=cache_dir, cache_name=cache_name, compress=compress)
+
+    def merge(self, other_grid: G, mode: str) -> G:
+        """Merge another grid into this grid. When ids overlap, ids of other_grid are offset to avoid conflicts.
+
+        Args:
+            other_grid (G): The grid to merge into this grid.
+            mode (str): The merge mode:
+                - "recalculate_ids": Recalculate ids of other_grid to avoid conflicts.
+                - "keep_ids": Keep ids of other_grid. Raises an error if grids contain overlapping indices.
+        """
+
+        merged_grid = merge_grids(self, other_grid, mode)
+
+        return merged_grid

tests/unit/model/grids/test_grid_base.py

@@ -15,6 +15,7 @@
     LineArray,
     LinkArray,
     NodeArray,
+    SourceArray,
     ThreeWindingTransformerArray,
     TransformerArray,
     TransformerTapRegulatorArray,
@@ -401,3 +402,42 @@ def test_from_txt_file(self, tmp_path: Path):
         assert 1 == grid.branches.filter(to_status=0).size
         assert 1 == grid.transformer.size
         np.testing.assert_array_equal([14, 10, 11, 12, 13, 15, 16, 17], grid.branches.id)
+
+
+class TestMergeGrids:
+    def test_merge_two_grids(self):
+        grid1 = Grid.from_txt("S1 2", "S1 3 link", "3 4 transformer")
+        grid2 = Grid.from_txt("S11 12", "S11 13 link", "13 14 transformer")
+        grid1_size = grid1.node.size
+        grid2_size = grid2.node.size
+
+        merged_grid = grid1.merge(grid2, mode="keep_ids")
+        merged_grid_size = merged_grid.node.size
+
+        assert merged_grid_size == grid1_size + grid2_size, "Merged grid size should be the sum of both grids' sizes"
+
+    def test_merge_two_grids_with_overlapping_node_ids(self):
+        grid1 = Grid.from_txt("S1 2", "S1 3 link", "3 4 transformer")
+        grid2 = Grid.from_txt("S1 2", "S1 13 link", "13 14 transformer")
+        source = SourceArray(id=[501], node=[1], status=[1], u_ref=[0.0])
+        grid1.append(source)
+        grid2.append(source)
+
+        merged_grid = grid1.merge(grid2, mode="recalculate_ids")
+        assert merged_grid.check_ids() is None, "Asset ids are not unique after merging!"
+
+        # Check if from and to nodes are updated by checking that their values form the entire set of node ids:
+        assert set(merged_grid.branches.from_node).union(merged_grid.branches.to_node) == set(merged_grid.node.id), (
+            "All from and to nodes should form the entire set of node ids in the merged grid!"
+        )
+
+        # assert node in grid.source is updated by checking if the node column contains values that are all node ids:
+        assert set(merged_grid.source.node).issubset(merged_grid.node.id), "All source nodes should be valid node ids!"
+
+    def test_merge_two_grids_with_overlapping_line(self):
+        # Now both grids have 14 as highest node id, so both will have branch ids 15, 16 and 17:
+        grid1 = Grid.from_txt("S1 2", "S1 3 link", "3 14 transformer")
+        grid2 = Grid.from_txt("S1 2", "S1 13 link", "13 14 transformer")
+
+        merged_grid = grid1.merge(grid2, mode="recalculate_ids")
+        assert merged_grid.check_ids() is None, "Asset ids are not unique after merging!"