Add allocation network id offsetting

Deltares · Dec 20, 2023 · c3ef344 · c3ef344
1 parent f3e0710
commit c3ef344
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 0 deletions.
diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
@@ -58,6 +58,17 @@ def translate_spacially(
         )
         return edge
 
+    def offset_allocation_network_ids(
+        self, offset_allocation_network_id: int, inplace: bool = True
+    ) -> "Edge":
+        if inplace:
+            edge = self
+        else:
+            edge = deepcopy(self)
+
+        edge.df.allocation_network_id += offset_allocation_network_id
+        return edge
+
     def get_where_edge_type(self, edge_type: str) -> NDArray[np.bool_]:
         return (self.df.edge_type == edge_type).to_numpy()
 

diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
@@ -76,6 +76,17 @@ def translate_spacially(
         )
         return node
 
+    def offset_allocation_network_ids(
+        self, offset_allocation_network_id: int, inplace: bool = True
+    ) -> "Node":
+        if inplace:
+            node = self
+        else:
+            node = deepcopy(self)
+
+        node.df.allocation_network_id += offset_allocation_network_id
+        return node
+
     def geometry_from_connectivity(
         self, from_id: Sequence[int], to_id: Sequence[int]
     ) -> NDArray[Any]:

diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
@@ -71,6 +71,18 @@ def translate_spacially(
         network.edge.translate_spacially(offset_spacial)
         return network
 
+    def offset_allocation_network_ids(
+        self, offset_allocation_network_id: int, inplace: bool = True
+    ) -> "Network":
+        if inplace:
+            network = self
+        else:
+            network = deepcopy(self)
+
+        network.node.offset_allocation_network_ids(offset_allocation_network_id)
+        network.edge.offset_allocation_network_ids(offset_allocation_network_id)
+        return network
+
     @classmethod
     def _load(cls, filepath: Path | None) -> dict[str, Any]:
         directory = context_file_loading.get().get("directory", None)
@@ -392,10 +404,14 @@ def reset_contextvar(self) -> "Model":
     def max_node_id(self) -> int:
         return self.network.node.df.index.max()
 
+    def max_allocation_network_id(self) -> int:
+        return self.network.node.df.allocation_network_id.max()
+
     def merge_model(
         self,
         model_added: "Model",
         offset_node_id: int | None = None,
+        offset_allocation_network_id: int | None = None,
         offset_spacial: tuple[float, float] = (0.0, 0.0),
         inplace: bool = True,
     ):
@@ -410,6 +426,17 @@ def merge_model(
             offset_spacial, inplace=False
         )
         min_offset_node_id = model.max_node_id()
+        min_offset_allocation_network_id = model.max_allocation_network_id()
+
+        if offset_allocation_network_id is None:
+            offset_allocation_network_id = min_offset_allocation_network_id
+        else:
+            assert (
+                offset_allocation_network_id >= min_offset_allocation_network_id
+            ), f"The allocation network ID offset must be at least the maximum allocation network ID of the main model ({min_offset_allocation_network_id}) to avoid conflicts."
+        nodes_added["network"].offset_allocation_network_ids(
+            offset_allocation_network_id
+        )
 
         if offset_node_id is None:
             offset_node_id = min_offset_node_id