From c3ef3444e9d5f9ffe15faa8335d3ed918d41b2b7 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 14:11:07 +0100
Subject: [PATCH] Add allocation network id offsetting

---
 python/ribasim/ribasim/geometry/edge.py | 11 ++++++++++
 python/ribasim/ribasim/geometry/node.py | 11 ++++++++++
 python/ribasim/ribasim/model.py         | 27 +++++++++++++++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index a5efcec14..83365ab36 100644
--- 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
index 3153068f4..1b2eeb99f 100644
--- 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
index 93440d37b..f73bf4c47 100644
--- 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