From cd6f5325e18c85e2eb94b2cd8eddc476171c74f3 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 11:57:50 +0100
Subject: [PATCH 01/20] Add merging of models in ribasim Python

---
 python/ribasim/ribasim/geometry/edge.py | 20 +++++++
 python/ribasim/ribasim/geometry/node.py | 17 ++++++
 python/ribasim/ribasim/input_base.py    | 70 +++++++++++++++++++++++++
 python/ribasim/ribasim/model.py         | 51 ++++++++++++++++++
 4 files changed, 158 insertions(+)

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index 64b0b50b6..a5efcec14 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -1,3 +1,4 @@
+from copy import deepcopy
 from typing import Any
 
 import matplotlib.pyplot as plt
@@ -9,6 +10,7 @@
 from numpy.typing import NDArray
 from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
+from shapely.geometry import LineString
 
 from ribasim.input_base import SpatialTableModel
 
@@ -38,6 +40,24 @@ class Edge(SpatialTableModel[EdgeSchema]):
         Table describing the flow connections.
     """
 
+    def translate_spacially(
+        self, offset_spacial: tuple[float, float], inplace: bool = True
+    ) -> "Edge":
+        if inplace:
+            edge = self
+        else:
+            edge = deepcopy(self)
+
+        edge.df.geometry = edge.df.geometry.apply(
+            lambda linestring: LineString(
+                [
+                    (point[0] + offset_spacial[0], point[1] + offset_spacial[1])
+                    for point in linestring.coords
+                ]
+            )
+        )
+        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 a4f9f552e..3153068f4 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from copy import deepcopy
 from typing import Any
 
 import matplotlib.pyplot as plt
@@ -9,6 +10,7 @@
 from numpy.typing import NDArray
 from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
+from shapely.geometry import Point
 
 from ribasim.input_base import SpatialTableModel
 
@@ -59,6 +61,21 @@ def node_ids_and_types(*nodes):
 
         return node_id, node_type
 
+    def translate_spacially(
+        self, offset_spacial: tuple[float, float], inplace: bool = True
+    ) -> "Node":
+        if inplace:
+            node = self
+        else:
+            node = deepcopy(self)
+
+        node.df.geometry = node.df.geometry.apply(
+            lambda point: Point(
+                point.x + offset_spacial[0], point.y + offset_spacial[1]
+            )
+        )
+        return node
+
     def geometry_from_connectivity(
         self, from_id: Sequence[int], to_id: Sequence[int]
     ) -> NDArray[Any]:
diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 42958ab75..594beee86 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -3,6 +3,7 @@
 from collections.abc import Callable, Generator
 from contextlib import closing
 from contextvars import ContextVar
+from copy import deepcopy
 from pathlib import Path
 from sqlite3 import Connection, connect
 from typing import (
@@ -202,6 +203,42 @@ def node_ids(self) -> set[int]:
 
         return node_ids
 
+    def offset_node_ids(self, offset_node_id: int) -> "TableModel":
+        copy = deepcopy(self)
+        df = copy.df
+        if copy.df is not None:
+            df.index += offset_node_id
+            for name_column in [
+                "node_id",
+                "from_node_id",
+                "to_node_id",
+                "listen_node_id",
+            ]:
+                if hasattr(df, name_column):
+                    df[name_column] += offset_node_id
+        return copy
+
+    def merge_table(
+        self, table_added: "TableModel", inplace: bool = True
+    ) -> "TableModel":
+        assert type(self) == type(
+            table_added
+        ), "Can only merge tables of the same type."
+
+        if inplace:
+            table = self
+        else:
+            table = deepcopy(self)
+
+        table.df = pd.concat(
+            [
+                table.df,
+                table_added.df,
+            ]
+        )
+
+        return table
+
     @classmethod
     def _load(cls, filepath: Path | None) -> dict[str, Any]:
         db = context_file_loading.get().get("database")
@@ -421,6 +458,39 @@ def node_ids_and_types(self) -> tuple[list[int], list[str]]:
         ids = self.node_ids()
         return list(ids), len(ids) * [self.get_input_type()]
 
+    def offset_node_ids(self, offset_node_id: int) -> "NodeModel":
+        node_copy = deepcopy(self)
+        for field in node_copy.fields():
+            attr = getattr(node_copy, field)
+            if isinstance(attr, TableModel):
+                table = attr
+                setattr(
+                    node_copy,
+                    field,
+                    table.offset_node_ids(offset_node_id),
+                )
+        return node_copy
+
+    def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeModel":
+        assert type(self) == type(node_added), "Can only merge nodes of the same type."
+
+        if inplace:
+            node = self
+        else:
+            node = deepcopy(self)
+
+        for field in node_added.fields():
+            attr = getattr(node_added, field)
+            if isinstance(attr, TableModel):
+                table_added = attr
+                table_node = getattr(node, field)
+                if table_added.df is not None:
+                    if table_node.df is not None:
+                        table_added = table_node.merge_table(table_added, inplace=False)
+
+                    setattr(node, field, table_added)
+        return node
+
     def _save(self, directory: DirectoryPath, input_dir: DirectoryPath, **kwargs):
         for field in self.fields():
             getattr(self, field)._save(
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 5278b8efd..93440d37b 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -1,5 +1,6 @@
 import datetime
 import shutil
+from copy import deepcopy
 from pathlib import Path
 from typing import Any
 
@@ -58,6 +59,18 @@ def n_nodes(self):
 
         return n
 
+    def translate_spacially(
+        self, offset_spacial: tuple[float, float], inplace: bool = True
+    ) -> "Network":
+        if inplace:
+            network = self
+        else:
+            network = deepcopy(self)
+
+        network.node.translate_spacially(offset_spacial)
+        network.edge.translate_spacially(offset_spacial)
+        return network
+
     @classmethod
     def _load(cls, filepath: Path | None) -> dict[str, Any]:
         directory = context_file_loading.get().get("directory", None)
@@ -376,6 +389,44 @@ def reset_contextvar(self) -> "Model":
         context_file_loading.set({})
         return self
 
+    def max_node_id(self) -> int:
+        return self.network.node.df.index.max()
+
+    def merge_model(
+        self,
+        model_added: "Model",
+        offset_node_id: int | None = None,
+        offset_spacial: tuple[float, float] = (0.0, 0.0),
+        inplace: bool = True,
+    ):
+        if inplace:
+            model = self
+        else:
+            model = deepcopy(self)
+
+        nodes_model = model.nodes()
+        nodes_added = model_added.nodes()
+        nodes_added["network"] = nodes_added["network"].translate_spacially(
+            offset_spacial, inplace=False
+        )
+        min_offset_node_id = model.max_node_id()
+
+        if offset_node_id is None:
+            offset_node_id = min_offset_node_id
+        else:
+            assert (
+                offset_node_id >= min_offset_node_id
+            ), f"The node id offset must be at least the maximum node ID of the main model ({min_offset_node_id}) to avoid conflicts."
+
+        for node_type, node_added in nodes_added.items():
+            node_added = node_added.offset_node_ids(offset_node_id)
+            if node_type in nodes_model:
+                node_added = nodes_model[node_type].merge_node(
+                    node_added, inplace=False
+                )
+
+            setattr(model, node_type, node_added)
+
     def plot_control_listen(self, ax):
         x_start, x_end = [], []
         y_start, y_end = [], []

From 8a9943083d0eda3062d8e8bd095ee8e528c3310e Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 12:58:40 +0100
Subject: [PATCH 02/20] Allow partial basin state input in Ribasim core

---
 core/src/bmi.jl         | 12 +++---------
 core/src/utils.jl       | 27 ++++++++++++++++++++++-----
 core/test/utils_test.jl | 24 +++++++++++++-----------
 3 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index 4d25344f0..f90d9e84d 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -73,15 +73,9 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
     end
     @debug "Read database into memory."
 
-    storage = if isempty(state)
-        # default to nearly empty basins, perhaps make required input
-        fill(1.0, n)
-    else
-        storages, errors = get_storages_from_levels(parameters.basin, state.level)
-        if errors
-            error("Encountered errors while parsing the initial levels of basins.")
-        end
-        storages
+    storage, errors = get_storages_from_levels(parameters.basin, state.node_id, state.level)
+    if errors
+        error("Encountered errors while parsing the initial levels of basins.")
     end
     @assert length(storage) == n "Basin / state length differs from number of Basins"
     # Integrals for PID control
diff --git a/core/src/utils.jl b/core/src/utils.jl
index 5040f786c..f8b983272 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -376,12 +376,29 @@ end
 """Compute the storages of the basins based on the water level of the basins."""
 function get_storages_from_levels(
     basin::Basin,
-    levels::Vector,
+    state_node_id::Vector{Int},
+    state_level::Vector{Float64},
 )::Tuple{Vector{Float64}, Bool}
-    storages = Float64[]
-
-    for (i, level) in enumerate(levels)
-        push!(storages, get_storage_from_level(basin, i, level))
+    (; node_id) = basin
+
+    storages = fill(1.0, length(node_id))
+    n_specified_states = length(state_node_id)
+
+    if n_specified_states > 0
+        basin_state_index = 1
+        basin_state_node_id = state_node_id[1]
+
+        for (i, id) in enumerate(node_id)
+            if basin_state_node_id == id.value
+                storages[i] =
+                    get_storage_from_level(basin, i, state_level[basin_state_index])
+                basin_state_index += 1
+                if basin_state_index > n_specified_states
+                    break
+                end
+                basin_state_node_id = state_node_id[basin_state_index]
+            end
+        end
     end
     return storages, any(isnan.(storages))
 end
diff --git a/core/test/utils_test.jl b/core/test/utils_test.jl
index ed2bc1a1e..b9076459f 100644
--- a/core/test/utils_test.jl
+++ b/core/test/utils_test.jl
@@ -110,24 +110,26 @@ end
     ]
     storage = Ribasim.profile_storage(level, area)
     basin = Ribasim.Basin(
-        Indices(Ribasim.NodeID[1]),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        [area],
-        [level],
-        [storage],
+        Indices(Ribasim.NodeID[1, 2]),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        [area, area],
+        [level, level],
+        [storage, storage],
         StructVector{Ribasim.BasinTimeV1}(undef, 0),
     )
 
     logger = TestLogger()
     with_logger(logger) do
-        storages, errors = Ribasim.get_storages_from_levels(basin, [-1.0])
+        storages, errors = Ribasim.get_storages_from_levels(basin, [1], [-1.0])
         @test isnan(storages[1])
         @test errors
+        # Storage for basin with unspecified level is set to 1.0
+        @test storages[2] == 1.0
     end
 
     @test length(logger.logs) == 1

From f3e07109ab046e27f513904fe4363fe1f0d2e16b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 13:23:37 +0100
Subject: [PATCH 03/20] column name fix

---
 python/ribasim/ribasim/input_base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 594beee86..8bfa30bd5 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -212,7 +212,7 @@ def offset_node_ids(self, offset_node_id: int) -> "TableModel":
                 "node_id",
                 "from_node_id",
                 "to_node_id",
-                "listen_node_id",
+                "listen_feature_id",
             ]:
                 if hasattr(df, name_column):
                     df[name_column] += offset_node_id

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 04/20] 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

From f7e81a141a45d70c2975df20b95bf1bf36e5e377 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 14:27:42 +0100
Subject: [PATCH 05/20] Add allocation network plotting

---
 python/ribasim/ribasim/geometry/node.py | 44 +++++++++++++++++++++++++
 python/ribasim/ribasim/model.py         |  6 ++--
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index 1b2eeb99f..fd552b9c0 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -10,6 +10,7 @@
 from numpy.typing import NDArray
 from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
+from scipy.spatial import ConvexHull
 from shapely.geometry import Point
 
 from ribasim.input_base import SpatialTableModel
@@ -158,6 +159,49 @@ def connectivity_from_geometry(
         to_id = node_index[edge_node_id[:, 1]].to_numpy()
         return from_id, to_id
 
+    def plot_allocation_networks(self, ax=None, zorder=None) -> Any:
+        if ax is None:
+            _, ax = plt.subplots()
+            ax.axis("off")
+
+        for allocation_subnetwork_id, df_subnetwork in self.df.groupby(
+            "allocation_network_id"
+        ):
+            points = np.stack(
+                (
+                    df_subnetwork.geometry.x.to_numpy(),
+                    df_subnetwork.geometry.y.to_numpy(),
+                ),
+                axis=1,
+            )
+            hull = ConvexHull(points)
+            bounding_polygon_vertices = points[hull.vertices, :]
+            hull_center = np.mean(bounding_polygon_vertices, axis=0)
+            bounding_polygon_vertices = (
+                bounding_polygon_vertices - hull_center
+            ) * 1.1 + hull_center
+
+            ax.fill(
+                bounding_polygon_vertices[:, 0],
+                bounding_polygon_vertices[:, 1],
+                linestyle=":",
+                facecolor="none",
+                linewidth=2,
+                edgecolor="gray",
+                zorder=zorder,
+            )
+            if allocation_subnetwork_id == 1:
+                text = "Main network"
+            else:
+                text = f"Subnetwork {allocation_subnetwork_id}"
+            ax.text(
+                *hull_center,
+                text,
+                horizontalalignment="center",
+                color="gray",
+                zorder=zorder,
+            )
+
     def plot(self, ax=None, zorder=None) -> Any:
         """
         Plot the nodes. Each node type is given a separate marker.
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index f73bf4c47..bf5fa17ac 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -505,9 +505,9 @@ def plot_control_listen(self, ax):
                 label="Listen edge" if i == 0 else None,
             )
 
-    def plot(self, ax=None) -> Any:
+    def plot(self, ax=None, indicate_subnetworks: bool = True) -> Any:
         """
-        Plot the nodes and edges of the model.
+        Plot the nodes, edges and allocation networks of the model.
 
         Parameters
         ----------
@@ -524,6 +524,8 @@ def plot(self, ax=None) -> Any:
         self.network.edge.plot(ax=ax, zorder=2)
         self.plot_control_listen(ax)
         self.network.node.plot(ax=ax, zorder=3)
+        if indicate_subnetworks:
+            self.network.node.plot_allocation_networks(ax=ax)
 
         ax.legend(loc="lower left", bbox_to_anchor=(1, 0.5))
 

From bacb44106e435d150d06b875264ca7e2488de0cb Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 14:47:13 +0100
Subject: [PATCH 06/20] Add test model which merges other testmodels

---
 python/ribasim/ribasim/model.py               |   2 +-
 .../ribasim_testmodels/__init__.py            |   4 +-
 .../ribasim_testmodels/allocation.py          | 111 +++++++++++++++++-
 3 files changed, 112 insertions(+), 5 deletions(-)

diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index bf5fa17ac..2ecd711c1 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -525,7 +525,7 @@ def plot(self, ax=None, indicate_subnetworks: bool = True) -> Any:
         self.plot_control_listen(ax)
         self.network.node.plot(ax=ax, zorder=3)
         if indicate_subnetworks:
-            self.network.node.plot_allocation_networks(ax=ax)
+            self.network.node.plot_allocation_networks(ax=ax, zorder=4)
 
         ax.legend(loc="lower left", bbox_to_anchor=(1, 0.5))
 
diff --git a/python/ribasim_testmodels/ribasim_testmodels/__init__.py b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
index 7de7eb347..652acc790 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/__init__.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
@@ -7,8 +7,9 @@
 import ribasim_testmodels
 from ribasim_testmodels.allocation import (
     allocation_example_model,
-    # looped_subnetwork_model,
     fractional_flow_subnetwork_model,
+    # looped_subnetwork_model,
+    main_network_with_subnetworks_model,
     minimal_subnetwork_model,
     subnetwork_model,
     user_model,
@@ -85,6 +86,7 @@
     # Disable until this issue is resolved:
     # https://github.com/Deltares/Ribasim/issues/692
     # "looped_subnetwork_model",
+    "main_network_with_subnetworks_model",
 ]
 
 # provide a mapping from model name to its constructor, so we can iterate over all models
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index c3416c4d4..4fa62a5e9 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -123,7 +123,9 @@ def user_model():
 
 
 def subnetwork_model():
-    """Create a user testmodel representing a subnetwork."""
+    """Create a user testmodel representing a subnetwork.
+    This model is merged into main_network_with_subnetworks_model.
+    """
 
     # Setup the nodes:
     xy = np.array(
@@ -283,7 +285,9 @@ def subnetwork_model():
 
 
 def looped_subnetwork_model():
-    """Create a user testmodel representing a subnetwork containing a loop in the topology."""
+    """Create a user testmodel representing a subnetwork containing a loop in the topology.
+    This model is merged into main_network_with_subnetworks_model.
+    """
     # Setup the nodes:
     xy = np.array(
         [
@@ -670,7 +674,9 @@ def minimal_subnetwork_model():
 
 
 def fractional_flow_subnetwork_model():
-    """Create a small subnetwork that contains fractional flow nodes."""
+    """Create a small subnetwork that contains fractional flow nodes.
+    This model is merged into main_network_with_subnetworks_model.
+    """
 
     xy = np.array(
         [
@@ -1065,3 +1071,102 @@ def allocation_example_model():
     )
 
     return model
+
+
+def main_network_with_subnetworks_model():
+    """Generate a model with a main network with connected subnetworks which are other test models."""
+
+    # Main network
+    n_basins = 5
+    n_nodes = 2 * n_basins
+    xy = np.array([(3 * i, (-1) ** (i + 1)) for i in range(n_nodes)])
+    node_xy = gpd.points_from_xy(x=xy[:, 0], y=xy[:, 1])
+
+    node_type = ["FlowBoundary", "Basin"] + (n_basins - 1) * [
+        "LinearResistance",
+        "Basin",
+    ]
+    node = ribasim.Node(
+        df=gpd.GeoDataFrame(
+            data={"type": node_type, "allocation_network_id": 1},
+            index=pd.Index(np.arange(len(xy)) + 1, name="fid"),
+            geometry=node_xy,
+            crs="EPSG:28992",
+        )
+    )
+
+    from_id = np.arange(1, n_nodes)
+    to_id = np.arange(2, n_nodes + 1)
+    allocation_network_id = len(from_id) * [None]
+    allocation_network_id[0] = 1
+    lines = node.geometry_from_connectivity(from_id, to_id)
+    edge = ribasim.Edge(
+        df=gpd.GeoDataFrame(
+            data={
+                "from_node_id": from_id,
+                "to_node_id": to_id,
+                "edge_type": len(from_id) * ["flow"],
+                "allocation_network_id": allocation_network_id,
+            },
+            geometry=lines,
+            crs="EPSG:28992",
+        )
+    )
+
+    # Setup the basins:
+    basin_node_ids = np.arange(2, n_nodes + 1, 2)
+    profile = pd.DataFrame(
+        data={
+            "node_id": np.repeat(basin_node_ids, [2] * len(basin_node_ids)),
+            "area": 1000.0,
+            "level": n_basins * [0.0, 1.0],
+        }
+    )
+
+    static = pd.DataFrame(
+        data={
+            "node_id": basin_node_ids,
+            "drainage": 0.0,
+            "potential_evaporation": 0.0,
+            "infiltration": 0.0,
+            "precipitation": 0.0,
+            "urban_runoff": 0.0,
+        }
+    )
+
+    state = pd.DataFrame(data={"node_id": basin_node_ids, "level": 1.0})
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
+
+    flow_boundary = ribasim.FlowBoundary(
+        static=pd.DataFrame(
+            data={
+                "node_id": [1],
+                "flow_rate": 1.0,
+            }
+        )
+    )
+
+    linear_resistance = ribasim.LinearResistance(
+        static=pd.DataFrame(
+            data={"node_id": np.arange(3, n_nodes + 1, 2), "resistance": 1e-3}
+        )
+    )
+
+    model = ribasim.Model(
+        network=ribasim.Network(
+            node=node,
+            edge=edge,
+        ),
+        basin=basin,
+        flow_boundary=flow_boundary,
+        linear_resistance=linear_resistance,
+        starttime="2020-01-01 00:00:00",
+        endtime="2021-01-01 00:00:00",
+    )
+
+    model.merge_model(subnetwork_model(), offset_spacial=(0.0, 3.0))
+    model.merge_model(fractional_flow_subnetwork_model(), offset_spacial=(14.0, 3.0))
+    model.merge_model(looped_subnetwork_model(), offset_spacial=(26.0, 3.0))
+
+    return model

From 7accfdc693c435accd14c17784529746df6e725d Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 16:09:40 +0100
Subject: [PATCH 07/20] Add method for adding edges to existing model

---
 python/ribasim/ribasim/geometry/node.py       |  7 ++-
 python/ribasim/ribasim/model.py               | 43 +++++++++++++++++++
 .../ribasim_testmodels/allocation.py          |  6 +++
 3 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index fd552b9c0..adc5f44d3 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -184,10 +184,9 @@ def plot_allocation_networks(self, ax=None, zorder=None) -> Any:
             ax.fill(
                 bounding_polygon_vertices[:, 0],
                 bounding_polygon_vertices[:, 1],
-                linestyle=":",
-                facecolor="none",
-                linewidth=2,
-                edgecolor="gray",
+                facecolor="gray",
+                alpha=0.5,
+                linewidth=0,
                 zorder=zorder,
             )
             if allocation_subnetwork_id == 1:
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 2ecd711c1..daee19ab4 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import Any
 
+import geopandas as gpd
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -17,6 +18,7 @@
     model_serializer,
     model_validator,
 )
+from shapely.geometry import LineString
 
 from ribasim.config import (
     Allocation,
@@ -83,6 +85,47 @@ def offset_allocation_network_ids(
         network.edge.offset_allocation_network_ids(offset_allocation_network_id)
         return network
 
+    def add_edges(
+        self,
+        from_node_id: np.ndarray[int],
+        to_node_id: np.ndarray[int],
+        edge_type: list[str],
+        inplace: bool = True,
+    ) -> "Network":
+        if inplace:
+            network = self
+        else:
+            network = deepcopy(self)
+
+        df = pd.DataFrame(
+            data={
+                "from_node_id": from_node_id,
+                "to_node_id": to_node_id,
+                "edge_type": edge_type,
+            }
+        )
+
+        df["geometry"] = df.apply(
+            (
+                lambda row: LineString(
+                    [
+                        self.node.df.loc[row.from_node_id].geometry,
+                        self.node.df.loc[row.to_node_id].geometry,
+                    ]
+                )
+            ),
+            axis=1,
+        )
+        edges_added = gpd.GeoDataFrame(df, geometry=df.geometry)
+
+        network.edge.df = pd.concat(
+            [
+                network.edge.df,
+                edges_added,
+            ]
+        )
+        return network
+
     @classmethod
     def _load(cls, filepath: Path | None) -> dict[str, Any]:
         directory = context_file_loading.get().get("directory", None)
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 4fa62a5e9..b9916e1c9 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1169,4 +1169,10 @@ def main_network_with_subnetworks_model():
     model.merge_model(fractional_flow_subnetwork_model(), offset_spacial=(14.0, 3.0))
     model.merge_model(looped_subnetwork_model(), offset_spacial=(26.0, 3.0))
 
+    # Connection edges
+    from_id = np.array([2, 6, 10])
+    to_id = np.array([11, 24, 38])
+    edge_type = 3 * ["flow"]
+    model.network.add_edges(from_id, to_id, edge_type)
+
     return model

From be96b1da5a1957a5be289cd10d8295e128dac2f1 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 20 Dec 2023 16:49:59 +0100
Subject: [PATCH 08/20] support deleting nodes by id

---
 python/ribasim/ribasim/input_base.py          | 20 +++++++++++++++++++
 python/ribasim/ribasim/model.py               |  4 +++-
 .../ribasim_testmodels/allocation.py          | 10 ++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 8bfa30bd5..4c61a26be 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -13,6 +13,7 @@
 )
 
 import geopandas as gpd
+import numpy as np
 import pandas as pd
 import pandera as pa
 from pandera.typing import DataFrame
@@ -491,6 +492,25 @@ def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeMode
                     setattr(node, field, table_added)
         return node
 
+    def delete_by_ids(
+        self, node_ids: np.ndarray[int], inplace: bool = True
+    ) -> "NodeModel":
+        if inplace:
+            node = self
+        else:
+            node = deepcopy(self)
+
+        for field in node.fields():
+            attr = getattr(node, field)
+            if isinstance(attr, TableModel):
+                df = attr.df[~attr.df.node_id.isin(node_ids)]
+                if df.empty:
+                    attr.df = None
+                else:
+                    attr.df = df
+
+        return node
+
     def _save(self, directory: DirectoryPath, input_dir: DirectoryPath, **kwargs):
         for field in self.fields():
             getattr(self, field)._save(
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index daee19ab4..47aeffeee 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -116,7 +116,9 @@ def add_edges(
             ),
             axis=1,
         )
-        edges_added = gpd.GeoDataFrame(df, geometry=df.geometry)
+        edges_added = gpd.GeoDataFrame(
+            df, geometry=df.geometry, crs=network.edge.df.crs
+        )
 
         network.edge.df = pd.concat(
             [
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index b9916e1c9..e92aee508 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1175,4 +1175,14 @@ def main_network_with_subnetworks_model():
     edge_type = 3 * ["flow"]
     model.network.add_edges(from_id, to_id, edge_type)
 
+    # Convert connecting flow boundaries to pumps
+    model.network.node.df.loc[to_id, "type"] = "Pump"
+    model.flow_boundary.delete_by_ids(to_id)
+    pump_added = ribasim.Pump(
+        static=pd.DataFrame(
+            data={"node_id": to_id, "flow_rate": 1e-3, "max_flow_rate": 1.0}
+        )
+    )
+    model.pump.merge_node(pump_added)
+
     return model

From 3b7b44c4b56c0d296b458227a816bf7e9f10d99b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 21 Dec 2023 11:21:09 +0100
Subject: [PATCH 09/20] Remove scipy dependency, add subnetworks to legend

---
 python/ribasim/ribasim/geometry/node.py | 64 ++++++++++++-------------
 python/ribasim/ribasim/model.py         | 13 ++++-
 2 files changed, 43 insertions(+), 34 deletions(-)

diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index adc5f44d3..e5c17a4ff 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -2,15 +2,16 @@
 from copy import deepcopy
 from typing import Any
 
+import geopandas as gpd
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import pandera as pa
 import shapely
+from matplotlib.patches import Patch
 from numpy.typing import NDArray
 from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
-from scipy.spatial import ConvexHull
 from shapely.geometry import Point
 
 from ribasim.input_base import SpatialTableModel
@@ -164,42 +165,41 @@ def plot_allocation_networks(self, ax=None, zorder=None) -> Any:
             _, ax = plt.subplots()
             ax.axis("off")
 
+        COLOR_SUBNETWORK = "black"
+        COLOR_MAIN_NETWORK = "blue"
+        ALPHA = 0.25
+
+        contains_main_network = False
+        contains_subnetworks = False
+
         for allocation_subnetwork_id, df_subnetwork in self.df.groupby(
             "allocation_network_id"
         ):
-            points = np.stack(
-                (
-                    df_subnetwork.geometry.x.to_numpy(),
-                    df_subnetwork.geometry.y.to_numpy(),
-                ),
-                axis=1,
-            )
-            hull = ConvexHull(points)
-            bounding_polygon_vertices = points[hull.vertices, :]
-            hull_center = np.mean(bounding_polygon_vertices, axis=0)
-            bounding_polygon_vertices = (
-                bounding_polygon_vertices - hull_center
-            ) * 1.1 + hull_center
-
-            ax.fill(
-                bounding_polygon_vertices[:, 0],
-                bounding_polygon_vertices[:, 1],
-                facecolor="gray",
-                alpha=0.5,
-                linewidth=0,
-                zorder=zorder,
-            )
-            if allocation_subnetwork_id == 1:
-                text = "Main network"
+            if allocation_subnetwork_id is None:
+                continue
+            elif allocation_subnetwork_id == 1:
+                contains_main_network = True
+                color = COLOR_MAIN_NETWORK
             else:
-                text = f"Subnetwork {allocation_subnetwork_id}"
-            ax.text(
-                *hull_center,
-                text,
-                horizontalalignment="center",
-                color="gray",
-                zorder=zorder,
+                contains_subnetworks = True
+                color = COLOR_SUBNETWORK
+
+            hull = gpd.GeoDataFrame(
+                geometry=[df_subnetwork.geometry.unary_union.convex_hull]
             )
+            hull.plot(ax=ax, color=color, alpha=ALPHA, zorder=zorder)
+
+        handles = []
+        labels = []
+
+        if contains_main_network:
+            handles.append(Patch(facecolor=COLOR_MAIN_NETWORK, alpha=ALPHA))
+            labels.append("Main network")
+        if contains_subnetworks:
+            handles.append(Patch(facecolor=COLOR_SUBNETWORK, alpha=ALPHA))
+            labels.append("Subnetwork")
+
+        return handles, labels
 
     def plot(self, ax=None, zorder=None) -> Any:
         """
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 47aeffeee..7b8fa1e60 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -566,13 +566,22 @@ def plot(self, ax=None, indicate_subnetworks: bool = True) -> Any:
         if ax is None:
             _, ax = plt.subplots()
             ax.axis("off")
+
         self.network.edge.plot(ax=ax, zorder=2)
         self.plot_control_listen(ax)
         self.network.node.plot(ax=ax, zorder=3)
+
+        handles, labels = ax.get_legend_handles_labels()
+
         if indicate_subnetworks:
-            self.network.node.plot_allocation_networks(ax=ax, zorder=4)
+            (
+                handles_subnetworks,
+                labels_subnetworks,
+            ) = self.network.node.plot_allocation_networks(ax=ax, zorder=1)
+            handles += handles_subnetworks
+            labels += labels_subnetworks
 
-        ax.legend(loc="lower left", bbox_to_anchor=(1, 0.5))
+        ax.legend(handles, labels, loc="lower left", bbox_to_anchor=(1, 0.5))
 
         return ax
 

From 042e59561cf5fc2bffba6ada837cbf81197c775b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 21 Dec 2023 13:30:29 +0100
Subject: [PATCH 10/20] Make MyPy somewhat happier

---
 python/ribasim/ribasim/input_base.py | 25 ++++++++++++-------------
 python/ribasim/ribasim/model.py      |  4 ++--
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 4c61a26be..21d359276 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -13,7 +13,6 @@
 )
 
 import geopandas as gpd
-import numpy as np
 import pandas as pd
 import pandera as pa
 from pandera.typing import DataFrame
@@ -204,10 +203,10 @@ def node_ids(self) -> set[int]:
 
         return node_ids
 
-    def offset_node_ids(self, offset_node_id: int) -> "TableModel":
+    def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]":
         copy = deepcopy(self)
         df = copy.df
-        if copy.df is not None:
+        if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)):
             df.index += offset_node_id
             for name_column in [
                 "node_id",
@@ -220,8 +219,8 @@ def offset_node_ids(self, offset_node_id: int) -> "TableModel":
         return copy
 
     def merge_table(
-        self, table_added: "TableModel", inplace: bool = True
-    ) -> "TableModel":
+        self, table_added: "TableModel[TableT]", inplace: bool = True
+    ) -> "TableModel[TableT]":
         assert type(self) == type(
             table_added
         ), "Can only merge tables of the same type."
@@ -492,9 +491,7 @@ def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeMode
                     setattr(node, field, table_added)
         return node
 
-    def delete_by_ids(
-        self, node_ids: np.ndarray[int], inplace: bool = True
-    ) -> "NodeModel":
+    def delete_by_ids(self, node_ids: list[int], inplace: bool = True) -> "NodeModel":
         if inplace:
             node = self
         else:
@@ -503,11 +500,13 @@ def delete_by_ids(
         for field in node.fields():
             attr = getattr(node, field)
             if isinstance(attr, TableModel):
-                df = attr.df[~attr.df.node_id.isin(node_ids)]
-                if df.empty:
-                    attr.df = None
-                else:
-                    attr.df = df
+                df = attr.df
+                if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)):
+                    df = df[~df.node_id.isin(node_ids)]
+                    if df.empty:
+                        attr.df = None
+                    else:
+                        attr.df = df
 
         return node
 
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 7b8fa1e60..1d0699a4d 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -87,8 +87,8 @@ def offset_allocation_network_ids(
 
     def add_edges(
         self,
-        from_node_id: np.ndarray[int],
-        to_node_id: np.ndarray[int],
+        from_node_id: list[int],
+        to_node_id: list[int],
         edge_type: list[str],
         inplace: bool = True,
     ) -> "Network":

From 553c956942d48856e01e52112add656ad0a7c04e Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 21 Dec 2023 19:45:48 +0100
Subject: [PATCH 11/20] Add tests

---
 python/ribasim/ribasim/input_base.py | 13 +++++++++++++
 python/ribasim/ribasim/model.py      |  5 ++++-
 python/ribasim/tests/conftest.py     |  5 +++++
 python/ribasim/tests/test_model.py   | 17 +++++++++++++++++
 4 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 21d359276..e0a89d1b2 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -160,6 +160,19 @@ def _load(cls, filepath: Path | None) -> dict[str, Any]:
 class TableModel(FileModel, Generic[TableT]):
     df: DataFrame[TableT] | None = Field(default=None, exclude=True, repr=False)
 
+    def __eq__(self, other) -> bool:
+        if not type(self) == type(other):
+            return False
+        if self.filepath != other.filepath:
+            return False
+
+        if self.df is None and other.df is None:
+            return True
+        elif isinstance(self.df, (pd.DataFrame, gpd.GeoDataFrame)):
+            return self.df.equals(other.df)
+        else:
+            return False
+
     @field_validator("df")
     @classmethod
     def prefix_extra_columns(cls, v: DataFrame[TableT]):
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 1d0699a4d..531ec22d9 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -450,7 +450,10 @@ 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()
+        m = self.network.node.df.allocation_network_id.max()
+        if pd.isna(m):
+            m = 0
+        return m
 
     def merge_model(
         self,
diff --git a/python/ribasim/tests/conftest.py b/python/ribasim/tests/conftest.py
index 4e4c1d09f..8efffb016 100644
--- a/python/ribasim/tests/conftest.py
+++ b/python/ribasim/tests/conftest.py
@@ -32,3 +32,8 @@ def backwater() -> ribasim.Model:
 @pytest.fixture()
 def discrete_control_of_pid_control() -> ribasim.Model:
     return ribasim_testmodels.discrete_control_of_pid_control_model()
+
+
+@pytest.fixture()
+def subnetwork() -> ribasim.Model:
+    return ribasim_testmodels.subnetwork_model()
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index 23da78e53..d4da03a3b 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -1,4 +1,5 @@
 import re
+from copy import deepcopy
 from sqlite3 import connect
 
 import pandas as pd
@@ -143,3 +144,19 @@ def test_write_adds_fid_in_tables(basic, tmp_path):
         assert "fid" in df.columns
         fids = df.get("fid")
         assert fids.equals(pd.Series(range(1, len(fids) + 1)))
+
+
+def test_model_merging(basic, subnetwork, tmp_path):
+    model = deepcopy(basic)
+    model_added = deepcopy(subnetwork)
+    model.merge_model(model_added)
+    model.merge_model(model_added)
+    assert (model.network.node.df.index == range(1, 44)).all()
+    assert model.max_allocation_network_id() == 2
+    for node_type, node_added in model_added.nodes().items():
+        node_subnetwork = getattr(subnetwork, node_type)
+        for table_added, table_subnetwork in zip(
+            node_added.tables(), node_subnetwork.tables()
+        ):
+            assert table_added == table_subnetwork
+    model.write(tmp_path / "compound_model/ribasim.toml")

From ff040aa2ba4417130ca934954688675a3a2a1a07 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 21 Dec 2023 19:54:29 +0100
Subject: [PATCH 12/20] Expand plot testing

---
 python/ribasim/tests/conftest.py   | 5 +++++
 python/ribasim/tests/test_model.py | 3 ++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/python/ribasim/tests/conftest.py b/python/ribasim/tests/conftest.py
index 8efffb016..611fc9c76 100644
--- a/python/ribasim/tests/conftest.py
+++ b/python/ribasim/tests/conftest.py
@@ -37,3 +37,8 @@ def discrete_control_of_pid_control() -> ribasim.Model:
 @pytest.fixture()
 def subnetwork() -> ribasim.Model:
     return ribasim_testmodels.subnetwork_model()
+
+
+@pytest.fixture()
+def main_network_with_subnetworks() -> ribasim.Model:
+    return ribasim_testmodels.main_network_with_subnetworks_model()
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index d4da03a3b..36998dc6f 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -131,8 +131,9 @@ def test_tabulated_rating_curve_model(tabulated_rating_curve, tmp_path):
     Model.read(tmp_path / "tabulated_rating_curve/ribasim.toml")
 
 
-def test_plot(discrete_control_of_pid_control):
+def test_plot(discrete_control_of_pid_control, main_network_with_subnetworks):
     discrete_control_of_pid_control.plot()
+    main_network_with_subnetworks.plot()
 
 
 def test_write_adds_fid_in_tables(basic, tmp_path):

From 6049ae3939b12d847e5cb508ddcfdbb1bd993d59 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Fri, 22 Dec 2023 17:00:53 +0100
Subject: [PATCH 13/20] Add docstrings

---
 python/ribasim/ribasim/geometry/edge.py | 2 ++
 python/ribasim/ribasim/geometry/node.py | 2 ++
 python/ribasim/ribasim/input_base.py    | 5 +++++
 python/ribasim/ribasim/model.py         | 8 ++++++++
 python/ribasim/tests/test_model.py      | 1 +
 5 files changed, 18 insertions(+)

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index 83365ab36..8d17f53c4 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -43,6 +43,7 @@ class Edge(SpatialTableModel[EdgeSchema]):
     def translate_spacially(
         self, offset_spacial: tuple[float, float], inplace: bool = True
     ) -> "Edge":
+        """Add the same spacial offset to all edges."""
         if inplace:
             edge = self
         else:
@@ -61,6 +62,7 @@ def translate_spacially(
     def offset_allocation_network_ids(
         self, offset_allocation_network_id: int, inplace: bool = True
     ) -> "Edge":
+        """Add the same offset to all node IDs."""
         if inplace:
             edge = self
         else:
diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index e5c17a4ff..a6ad869a6 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -66,6 +66,7 @@ def node_ids_and_types(*nodes):
     def translate_spacially(
         self, offset_spacial: tuple[float, float], inplace: bool = True
     ) -> "Node":
+        """Add the same spacial offset to all nodes."""
         if inplace:
             node = self
         else:
@@ -81,6 +82,7 @@ def translate_spacially(
     def offset_allocation_network_ids(
         self, offset_allocation_network_id: int, inplace: bool = True
     ) -> "Node":
+        """Add the same offset to all node IDs."""
         if inplace:
             node = self
         else:
diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index e0a89d1b2..474ab9946 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -217,6 +217,7 @@ def node_ids(self) -> set[int]:
         return node_ids
 
     def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]":
+        """Add the same offset to all node IDs."""
         copy = deepcopy(self)
         df = copy.df
         if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)):
@@ -234,6 +235,7 @@ def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]":
     def merge_table(
         self, table_added: "TableModel[TableT]", inplace: bool = True
     ) -> "TableModel[TableT]":
+        """Merge an added table of the same type into this table."""
         assert type(self) == type(
             table_added
         ), "Can only merge tables of the same type."
@@ -472,6 +474,7 @@ def node_ids_and_types(self) -> tuple[list[int], list[str]]:
         return list(ids), len(ids) * [self.get_input_type()]
 
     def offset_node_ids(self, offset_node_id: int) -> "NodeModel":
+        """Add the same offset to all node IDs in all underlying tables."""
         node_copy = deepcopy(self)
         for field in node_copy.fields():
             attr = getattr(node_copy, field)
@@ -485,6 +488,7 @@ def offset_node_ids(self, offset_node_id: int) -> "NodeModel":
         return node_copy
 
     def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeModel":
+        """Merge an added node of the same type into this node."""
         assert type(self) == type(node_added), "Can only merge nodes of the same type."
 
         if inplace:
@@ -505,6 +509,7 @@ def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeMode
         return node
 
     def delete_by_ids(self, node_ids: list[int], inplace: bool = True) -> "NodeModel":
+        """Delete all rows of the underlying tables whose node ID is in the given list."""
         if inplace:
             node = self
         else:
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 531ec22d9..e3c0eccbc 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -64,6 +64,7 @@ def n_nodes(self):
     def translate_spacially(
         self, offset_spacial: tuple[float, float], inplace: bool = True
     ) -> "Network":
+        """Add the same spacial offset to all nodes and edges."""
         if inplace:
             network = self
         else:
@@ -92,6 +93,7 @@ def add_edges(
         edge_type: list[str],
         inplace: bool = True,
     ) -> "Network":
+        """Add new edges to the network of the given type. Assumes no source edges are added."""
         if inplace:
             network = self
         else:
@@ -463,6 +465,12 @@ def merge_model(
         offset_spacial: tuple[float, float] = (0.0, 0.0),
         inplace: bool = True,
     ):
+        """
+        Merge copies of the nodes and edges of an added model into this model.
+        The added model is not modified, but the following modificadions are made to the added data:
+        - Node IDs are shifted by at least the maximum node ID of this model
+        - Allocation network IDs are shifted by at least the maximum allocation network ID of this model.
+        """
         if inplace:
             model = self
         else:
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index 36998dc6f..9ecb2fbfb 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -154,6 +154,7 @@ def test_model_merging(basic, subnetwork, tmp_path):
     model.merge_model(model_added)
     assert (model.network.node.df.index == range(1, 44)).all()
     assert model.max_allocation_network_id() == 2
+    # Added model should not change
     for node_type, node_added in model_added.nodes().items():
         node_subnetwork = getattr(subnetwork, node_type)
         for table_added, table_subnetwork in zip(

From e4976d18d470056babf9350129bda08eef4ac2dd Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Fri, 22 Dec 2023 17:11:21 +0100
Subject: [PATCH 14/20] Make test model runnable

---
 python/ribasim_testmodels/ribasim_testmodels/allocation.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index e92aee508..387c9ce80 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1185,4 +1185,9 @@ def main_network_with_subnetworks_model():
     )
     model.pump.merge_node(pump_added)
 
+    # Fix control logic
+    df = model.discrete_control.condition.df
+    df.listen_feature_id = 25
+    df.variable = "level"
+
     return model

From f450f72416eddd69d68e277feec45fa41db46acc Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 8 Jan 2024 16:10:06 +0100
Subject: [PATCH 15/20] spacial -> spatial :D

---
 python/ribasim/ribasim/geometry/edge.py          |  8 ++++----
 python/ribasim/ribasim/geometry/node.py          |  8 ++++----
 python/ribasim/ribasim/model.py                  | 16 ++++++++--------
 .../ribasim_testmodels/allocation.py             |  6 +++---
 4 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index 8d17f53c4..a7a7155bb 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -40,10 +40,10 @@ class Edge(SpatialTableModel[EdgeSchema]):
         Table describing the flow connections.
     """
 
-    def translate_spacially(
-        self, offset_spacial: tuple[float, float], inplace: bool = True
+    def translate_spatially(
+        self, offset_spatial: tuple[float, float], inplace: bool = True
     ) -> "Edge":
-        """Add the same spacial offset to all edges."""
+        """Add the same spatial offset to all edges."""
         if inplace:
             edge = self
         else:
@@ -52,7 +52,7 @@ def translate_spacially(
         edge.df.geometry = edge.df.geometry.apply(
             lambda linestring: LineString(
                 [
-                    (point[0] + offset_spacial[0], point[1] + offset_spacial[1])
+                    (point[0] + offset_spatial[0], point[1] + offset_spatial[1])
                     for point in linestring.coords
                 ]
             )
diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index a6ad869a6..dcf32333f 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -63,10 +63,10 @@ def node_ids_and_types(*nodes):
 
         return node_id, node_type
 
-    def translate_spacially(
-        self, offset_spacial: tuple[float, float], inplace: bool = True
+    def translate_spatially(
+        self, offset_spatial: tuple[float, float], inplace: bool = True
     ) -> "Node":
-        """Add the same spacial offset to all nodes."""
+        """Add the same spatial offset to all nodes."""
         if inplace:
             node = self
         else:
@@ -74,7 +74,7 @@ def translate_spacially(
 
         node.df.geometry = node.df.geometry.apply(
             lambda point: Point(
-                point.x + offset_spacial[0], point.y + offset_spacial[1]
+                point.x + offset_spatial[0], point.y + offset_spatial[1]
             )
         )
         return node
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index e3c0eccbc..c6a90aeda 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -61,17 +61,17 @@ def n_nodes(self):
 
         return n
 
-    def translate_spacially(
-        self, offset_spacial: tuple[float, float], inplace: bool = True
+    def translate_spatially(
+        self, offset_spatial: tuple[float, float], inplace: bool = True
     ) -> "Network":
-        """Add the same spacial offset to all nodes and edges."""
+        """Add the same spatial offset to all nodes and edges."""
         if inplace:
             network = self
         else:
             network = deepcopy(self)
 
-        network.node.translate_spacially(offset_spacial)
-        network.edge.translate_spacially(offset_spacial)
+        network.node.translate_spatially(offset_spatial)
+        network.edge.translate_spatially(offset_spatial)
         return network
 
     def offset_allocation_network_ids(
@@ -462,7 +462,7 @@ def merge_model(
         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),
+        offset_spatial: tuple[float, float] = (0.0, 0.0),
         inplace: bool = True,
     ):
         """
@@ -478,8 +478,8 @@ def merge_model(
 
         nodes_model = model.nodes()
         nodes_added = model_added.nodes()
-        nodes_added["network"] = nodes_added["network"].translate_spacially(
-            offset_spacial, inplace=False
+        nodes_added["network"] = nodes_added["network"].translate_spatially(
+            offset_spatial, inplace=False
         )
         min_offset_node_id = model.max_node_id()
         min_offset_allocation_network_id = model.max_allocation_network_id()
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 387c9ce80..fc6c90df2 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1165,9 +1165,9 @@ def main_network_with_subnetworks_model():
         endtime="2021-01-01 00:00:00",
     )
 
-    model.merge_model(subnetwork_model(), offset_spacial=(0.0, 3.0))
-    model.merge_model(fractional_flow_subnetwork_model(), offset_spacial=(14.0, 3.0))
-    model.merge_model(looped_subnetwork_model(), offset_spacial=(26.0, 3.0))
+    model.merge_model(subnetwork_model(), offset_spatial=(0.0, 3.0))
+    model.merge_model(fractional_flow_subnetwork_model(), offset_spatial=(14.0, 3.0))
+    model.merge_model(looped_subnetwork_model(), offset_spatial=(26.0, 3.0))
 
     # Connection edges
     from_id = np.array([2, 6, 10])

From 8fe09499a1edf8db04058385d08003764eb41898 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 9 Jan 2024 10:53:56 +0100
Subject: [PATCH 16/20] Process some comments

---
 python/ribasim/ribasim/geometry/edge.py                | 10 +---------
 python/ribasim/ribasim/model.py                        |  4 ++--
 python/ribasim/tests/test_model.py                     |  4 ++--
 .../ribasim_testmodels/allocation.py                   |  6 +++---
 4 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index a7a7155bb..835695682 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -10,7 +10,6 @@
 from numpy.typing import NDArray
 from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
-from shapely.geometry import LineString
 
 from ribasim.input_base import SpatialTableModel
 
@@ -49,14 +48,7 @@ def translate_spatially(
         else:
             edge = deepcopy(self)
 
-        edge.df.geometry = edge.df.geometry.apply(
-            lambda linestring: LineString(
-                [
-                    (point[0] + offset_spatial[0], point[1] + offset_spatial[1])
-                    for point in linestring.coords
-                ]
-            )
-        )
+        edge.df.geometry = edge.df.geometry.translate(*offset_spatial)
         return edge
 
     def offset_allocation_network_ids(
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index c6a90aeda..d1c78321e 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -457,7 +457,7 @@ def max_allocation_network_id(self) -> int:
             m = 0
         return m
 
-    def merge_model(
+    def merge(
         self,
         model_added: "Model",
         offset_node_id: int | None = None,
@@ -467,7 +467,7 @@ def merge_model(
     ):
         """
         Merge copies of the nodes and edges of an added model into this model.
-        The added model is not modified, but the following modificadions are made to the added data:
+        The added model is not modified, but the following modifications are made to the added data:
         - Node IDs are shifted by at least the maximum node ID of this model
         - Allocation network IDs are shifted by at least the maximum allocation network ID of this model.
         """
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index 9ecb2fbfb..b851bb68a 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -150,8 +150,8 @@ def test_write_adds_fid_in_tables(basic, tmp_path):
 def test_model_merging(basic, subnetwork, tmp_path):
     model = deepcopy(basic)
     model_added = deepcopy(subnetwork)
-    model.merge_model(model_added)
-    model.merge_model(model_added)
+    model.merge(model_added)
+    model.merge(model_added)
     assert (model.network.node.df.index == range(1, 44)).all()
     assert model.max_allocation_network_id() == 2
     # Added model should not change
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index fc6c90df2..c4c2570d3 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1165,9 +1165,9 @@ def main_network_with_subnetworks_model():
         endtime="2021-01-01 00:00:00",
     )
 
-    model.merge_model(subnetwork_model(), offset_spatial=(0.0, 3.0))
-    model.merge_model(fractional_flow_subnetwork_model(), offset_spatial=(14.0, 3.0))
-    model.merge_model(looped_subnetwork_model(), offset_spatial=(26.0, 3.0))
+    model.merge(subnetwork_model(), offset_spatial=(0.0, 3.0))
+    model.merge(fractional_flow_subnetwork_model(), offset_spatial=(14.0, 3.0))
+    model.merge(looped_subnetwork_model(), offset_spatial=(26.0, 3.0))
 
     # Connection edges
     from_id = np.array([2, 6, 10])

From cfae794bb7f8a67974ac0245af8050cd631ffe81 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 9 Jan 2024 11:43:46 +0100
Subject: [PATCH 17/20] Test model translation

---
 python/ribasim/tests/test_model.py | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index b851bb68a..e41552a04 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -150,10 +150,38 @@ def test_write_adds_fid_in_tables(basic, tmp_path):
 def test_model_merging(basic, subnetwork, tmp_path):
     model = deepcopy(basic)
     model_added = deepcopy(subnetwork)
+    # Add model twice to test adding multiple subnetworks
     model.merge(model_added)
-    model.merge(model_added)
+    model.merge(model_added, offset_spatial=(1.0, 1.0))
     assert (model.network.node.df.index == range(1, 44)).all()
     assert model.max_allocation_network_id() == 2
+    n_nodes_basic = len(basic.network.node.df)
+    n_edges_basic = len(basic.network.edge.df)
+    n_nodes_subnetwork = len(subnetwork.network.node.df)
+    n_edges_subnetwork = len(subnetwork.network.edge.df)
+    node_geometry = model.network.node.df.geometry
+    edge_geometry = model.network.edge.df.geometry
+    assert (
+        node_geometry[n_nodes_basic : n_nodes_basic + n_nodes_subnetwork]
+        == subnetwork.network.node.df.geometry.to_numpy()
+    ).all()
+    assert (
+        node_geometry[
+            n_nodes_basic + n_nodes_subnetwork : n_nodes_basic + 2 * n_nodes_subnetwork
+        ]
+        == subnetwork.network.node.df.geometry.translate(1.0, 1.0).to_numpy()
+    ).all()
+
+    assert (
+        edge_geometry[n_edges_basic : n_edges_basic + n_edges_subnetwork]
+        == subnetwork.network.edge.df.geometry.to_numpy()
+    ).all()
+    assert (
+        edge_geometry[
+            n_edges_basic + n_edges_subnetwork : n_edges_basic + 2 * n_edges_subnetwork
+        ]
+        == subnetwork.network.edge.df.geometry.translate(1.0, 1.0).to_numpy()
+    ).all()
     # Added model should not change
     for node_type, node_added in model_added.nodes().items():
         node_subnetwork = getattr(subnetwork, node_type)

From 6e8bdae4312d3a8506638c22924b965a4d02ef91 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 9 Jan 2024 15:02:32 +0100
Subject: [PATCH 18/20] rename merge methods

---
 python/ribasim/ribasim/input_base.py                       | 6 +++---
 python/ribasim/ribasim/model.py                            | 4 +---
 python/ribasim_testmodels/ribasim_testmodels/allocation.py | 2 +-
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 474ab9946..671b84269 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -232,7 +232,7 @@ def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]":
                     df[name_column] += offset_node_id
         return copy
 
-    def merge_table(
+    def merge(
         self, table_added: "TableModel[TableT]", inplace: bool = True
     ) -> "TableModel[TableT]":
         """Merge an added table of the same type into this table."""
@@ -487,7 +487,7 @@ def offset_node_ids(self, offset_node_id: int) -> "NodeModel":
                 )
         return node_copy
 
-    def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeModel":
+    def merge(self, node_added: "NodeModel", inplace: bool = True) -> "NodeModel":
         """Merge an added node of the same type into this node."""
         assert type(self) == type(node_added), "Can only merge nodes of the same type."
 
@@ -503,7 +503,7 @@ def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeMode
                 table_node = getattr(node, field)
                 if table_added.df is not None:
                     if table_node.df is not None:
-                        table_added = table_node.merge_table(table_added, inplace=False)
+                        table_added = table_node.merge(table_added, inplace=False)
 
                     setattr(node, field, table_added)
         return node
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index d1c78321e..f8680eeb0 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -504,9 +504,7 @@ def merge(
         for node_type, node_added in nodes_added.items():
             node_added = node_added.offset_node_ids(offset_node_id)
             if node_type in nodes_model:
-                node_added = nodes_model[node_type].merge_node(
-                    node_added, inplace=False
-                )
+                node_added = nodes_model[node_type].merge(node_added, inplace=False)
 
             setattr(model, node_type, node_added)
 
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index c4c2570d3..202323e7e 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1183,7 +1183,7 @@ def main_network_with_subnetworks_model():
             data={"node_id": to_id, "flow_rate": 1e-3, "max_flow_rate": 1.0}
         )
     )
-    model.pump.merge_node(pump_added)
+    model.pump.merge(pump_added)
 
     # Fix control logic
     df = model.discrete_control.condition.df

From 7db5ff498794170261b835b74f35844e76a84151 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Jan 2024 13:11:50 +0100
Subject: [PATCH 19/20] Comments adressed

---
 python/ribasim/ribasim/geometry/edge.py       | 12 ---
 python/ribasim/ribasim/geometry/node.py       | 17 ----
 python/ribasim/ribasim/input_base.py          | 84 ++++++++++++++-----
 python/ribasim/ribasim/model.py               | 66 ++++-----------
 python/ribasim/tests/test_model.py            | 33 +-------
 .../ribasim_testmodels/allocation.py          | 16 +++-
 .../ribasim_testmodels/utils.py               |  6 ++
 7 files changed, 102 insertions(+), 132 deletions(-)
 create mode 100644 python/ribasim_testmodels/ribasim_testmodels/utils.py

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index 835695682..bbc786447 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -39,18 +39,6 @@ class Edge(SpatialTableModel[EdgeSchema]):
         Table describing the flow connections.
     """
 
-    def translate_spatially(
-        self, offset_spatial: tuple[float, float], inplace: bool = True
-    ) -> "Edge":
-        """Add the same spatial offset to all edges."""
-        if inplace:
-            edge = self
-        else:
-            edge = deepcopy(self)
-
-        edge.df.geometry = edge.df.geometry.translate(*offset_spatial)
-        return edge
-
     def offset_allocation_network_ids(
         self, offset_allocation_network_id: int, inplace: bool = True
     ) -> "Edge":
diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index dcf32333f..b87b1ff62 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -12,7 +12,6 @@
 from numpy.typing import NDArray
 from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
-from shapely.geometry import Point
 
 from ribasim.input_base import SpatialTableModel
 
@@ -63,22 +62,6 @@ def node_ids_and_types(*nodes):
 
         return node_id, node_type
 
-    def translate_spatially(
-        self, offset_spatial: tuple[float, float], inplace: bool = True
-    ) -> "Node":
-        """Add the same spatial offset to all nodes."""
-        if inplace:
-            node = self
-        else:
-            node = deepcopy(self)
-
-        node.df.geometry = node.df.geometry.apply(
-            lambda point: Point(
-                point.x + offset_spatial[0], point.y + offset_spatial[1]
-            )
-        )
-        return node
-
     def offset_allocation_network_ids(
         self, offset_allocation_network_id: int, inplace: bool = True
     ) -> "Node":
diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 671b84269..5dc65cc25 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -13,6 +13,7 @@
 )
 
 import geopandas as gpd
+import numpy as np
 import pandas as pd
 import pandera as pa
 from pandera.typing import DataFrame
@@ -216,10 +217,15 @@ def node_ids(self) -> set[int]:
 
         return node_ids
 
-    def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]":
+    def offset_node_ids(
+        self, offset_node_id: int, inplace: bool = True
+    ) -> "TableModel[TableT]":
         """Add the same offset to all node IDs."""
-        copy = deepcopy(self)
-        df = copy.df
+        if inplace:
+            node = self
+        else:
+            node = deepcopy(self)
+        df = node.df
         if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)):
             df.index += offset_node_id
             for name_column in [
@@ -230,7 +236,22 @@ def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]":
             ]:
                 if hasattr(df, name_column):
                     df[name_column] += offset_node_id
-        return copy
+
+        return node
+
+    def offset_edge_ids(
+        self, offset_edge_id: int, inplace=True
+    ) -> "TableModel[TableT]":
+        if self.tablename() == "Edge" and isinstance(self.df, gpd.GeoDataFrame):
+            if inplace:
+                edge = self
+            else:
+                edge = deepcopy(self)
+            df = edge.df
+            if isinstance(df, gpd.GeoDataFrame):
+                df.index += offset_edge_id
+
+        return edge
 
     def merge(
         self, table_added: "TableModel[TableT]", inplace: bool = True
@@ -245,12 +266,29 @@ def merge(
         else:
             table = deepcopy(self)
 
-        table.df = pd.concat(
-            [
-                table.df,
-                table_added.df,
-            ]
-        )
+        if table_added.df is not None and table.df is not None:
+            if isinstance(self.df, gpd.GeoDataFrame):
+                common_node_ids = np.intersect1d(
+                    table.df.index.to_numpy(), table_added.df.index.to_numpy()
+                )
+            else:
+                common_node_ids = np.intersect1d(
+                    table.df.node_id.to_numpy(), table_added.df.node_id.to_numpy()
+                )
+
+            assert (
+                common_node_ids.size == 0
+            ), f"Self and added table (of type {type(self)}) have common IDs: {common_node_ids}."
+
+            table.df = pd.concat(
+                [
+                    table.df,
+                    table_added.df,
+                ]
+            )  # type: ignore
+
+        elif table_added.df is not None:
+            table.df = table_added.df
 
         return table
 
@@ -463,7 +501,7 @@ def tables(self) -> Generator[TableModel[Any], Any, None]:
             if isinstance(attr, TableModel):
                 yield attr
 
-    def node_ids(self):
+    def node_ids(self) -> set[int]:
         node_ids: set[int] = set()
         for table in self.tables():
             node_ids.update(table.node_ids())
@@ -473,19 +511,22 @@ def node_ids_and_types(self) -> tuple[list[int], list[str]]:
         ids = self.node_ids()
         return list(ids), len(ids) * [self.get_input_type()]
 
-    def offset_node_ids(self, offset_node_id: int) -> "NodeModel":
+    def offset_node_ids(self, offset_node_id: int, inplace: bool = True) -> "NodeModel":
         """Add the same offset to all node IDs in all underlying tables."""
-        node_copy = deepcopy(self)
-        for field in node_copy.fields():
-            attr = getattr(node_copy, field)
+        if inplace:
+            node = self
+        else:
+            node = deepcopy(self)
+        for field in node.fields():
+            attr = getattr(node, field)
             if isinstance(attr, TableModel):
                 table = attr
                 setattr(
-                    node_copy,
+                    node,
                     field,
                     table.offset_node_ids(offset_node_id),
                 )
-        return node_copy
+        return node
 
     def merge(self, node_added: "NodeModel", inplace: bool = True) -> "NodeModel":
         """Merge an added node of the same type into this node."""
@@ -501,11 +542,8 @@ def merge(self, node_added: "NodeModel", inplace: bool = True) -> "NodeModel":
             if isinstance(attr, TableModel):
                 table_added = attr
                 table_node = getattr(node, field)
-                if table_added.df is not None:
-                    if table_node.df is not None:
-                        table_added = table_node.merge(table_added, inplace=False)
-
-                    setattr(node, field, table_added)
+                table_added = table_node.merge(table_added, inplace=False)
+                setattr(node, field, table_added)
         return node
 
     def delete_by_ids(self, node_ids: list[int], inplace: bool = True) -> "NodeModel":
@@ -520,7 +558,7 @@ def delete_by_ids(self, node_ids: list[int], inplace: bool = True) -> "NodeModel
             if isinstance(attr, TableModel):
                 df = attr.df
                 if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)):
-                    df = df[~df.node_id.isin(node_ids)]
+                    df = df[~df.node_id.isin(node_ids)]  # type: ignore
                     if df.empty:
                         attr.df = None
                     else:
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index f8680eeb0..a75b7cb5a 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -61,19 +61,6 @@ def n_nodes(self):
 
         return n
 
-    def translate_spatially(
-        self, offset_spatial: tuple[float, float], inplace: bool = True
-    ) -> "Network":
-        """Add the same spatial offset to all nodes and edges."""
-        if inplace:
-            network = self
-        else:
-            network = deepcopy(self)
-
-        network.node.translate_spatially(offset_spatial)
-        network.edge.translate_spatially(offset_spatial)
-        return network
-
     def offset_allocation_network_ids(
         self, offset_allocation_network_id: int, inplace: bool = True
     ) -> "Network":
@@ -127,7 +114,7 @@ def add_edges(
                 network.edge.df,
                 edges_added,
             ]
-        )
+        )  # type: ignore
         return network
 
     @classmethod
@@ -451,62 +438,45 @@ def reset_contextvar(self) -> "Model":
     def max_node_id(self) -> int:
         return self.network.node.df.index.max()
 
+    def max_edge_id(self) -> int:
+        return self.network.edge.df.index.max()
+
     def max_allocation_network_id(self) -> int:
         m = self.network.node.df.allocation_network_id.max()
         if pd.isna(m):
             m = 0
         return m
 
-    def merge(
+    def smart_merge(
         self,
         model_added: "Model",
-        offset_node_id: int | None = None,
-        offset_allocation_network_id: int | None = None,
-        offset_spatial: tuple[float, float] = (0.0, 0.0),
-        inplace: bool = True,
     ):
         """
         Merge copies of the nodes and edges of an added model into this model.
         The added model is not modified, but the following modifications are made to the added data:
-        - Node IDs are shifted by at least the maximum node ID of this model
-        - Allocation network IDs are shifted by at least the maximum allocation network ID of this model.
+        - Node IDs are shifted by the maximum node ID of this model
+        - Allocation network IDs are shifted by the maximum allocation network ID of this model
         """
-        if inplace:
-            model = self
-        else:
-            model = deepcopy(self)
 
-        nodes_model = model.nodes()
+        nodes_model = self.nodes()
         nodes_added = model_added.nodes()
-        nodes_added["network"] = nodes_added["network"].translate_spatially(
-            offset_spatial, 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
-        )
+        offset_node_id = self.max_node_id()
+        offset_edge_id = self.max_edge_id()
+        offset_allocation_network_id = self.max_allocation_network_id()
 
-        if offset_node_id is None:
-            offset_node_id = min_offset_node_id
-        else:
-            assert (
-                offset_node_id >= min_offset_node_id
-            ), f"The node id offset must be at least the maximum node ID of the main model ({min_offset_node_id}) to avoid conflicts."
+        network_added = nodes_added["network"].offset_allocation_network_ids(
+            offset_allocation_network_id, inplace=False
+        )
+        network_added.edge.offset_edge_ids(offset_edge_id)
+        nodes_added["network"] = network_added
 
         for node_type, node_added in nodes_added.items():
-            node_added = node_added.offset_node_ids(offset_node_id)
+            node_added = node_added.offset_node_ids(offset_node_id, inplace=False)
             if node_type in nodes_model:
                 node_added = nodes_model[node_type].merge(node_added, inplace=False)
 
-            setattr(model, node_type, node_added)
+            setattr(self, node_type, node_added)
 
     def plot_control_listen(self, ax):
         x_start, x_end = [], []
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index e41552a04..27243a171 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -151,37 +151,12 @@ def test_model_merging(basic, subnetwork, tmp_path):
     model = deepcopy(basic)
     model_added = deepcopy(subnetwork)
     # Add model twice to test adding multiple subnetworks
-    model.merge(model_added)
-    model.merge(model_added, offset_spatial=(1.0, 1.0))
+    model.smart_merge(model_added)
+    model.smart_merge(model_added)
     assert (model.network.node.df.index == range(1, 44)).all()
+    assert not model.network.edge.df.index.duplicated().any()
     assert model.max_allocation_network_id() == 2
-    n_nodes_basic = len(basic.network.node.df)
-    n_edges_basic = len(basic.network.edge.df)
-    n_nodes_subnetwork = len(subnetwork.network.node.df)
-    n_edges_subnetwork = len(subnetwork.network.edge.df)
-    node_geometry = model.network.node.df.geometry
-    edge_geometry = model.network.edge.df.geometry
-    assert (
-        node_geometry[n_nodes_basic : n_nodes_basic + n_nodes_subnetwork]
-        == subnetwork.network.node.df.geometry.to_numpy()
-    ).all()
-    assert (
-        node_geometry[
-            n_nodes_basic + n_nodes_subnetwork : n_nodes_basic + 2 * n_nodes_subnetwork
-        ]
-        == subnetwork.network.node.df.geometry.translate(1.0, 1.0).to_numpy()
-    ).all()
-
-    assert (
-        edge_geometry[n_edges_basic : n_edges_basic + n_edges_subnetwork]
-        == subnetwork.network.edge.df.geometry.to_numpy()
-    ).all()
-    assert (
-        edge_geometry[
-            n_edges_basic + n_edges_subnetwork : n_edges_basic + 2 * n_edges_subnetwork
-        ]
-        == subnetwork.network.edge.df.geometry.translate(1.0, 1.0).to_numpy()
-    ).all()
+
     # Added model should not change
     for node_type, node_added in model_added.nodes().items():
         node_subnetwork = getattr(subnetwork, node_type)
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 202323e7e..aa4a78eae 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -3,6 +3,8 @@
 import pandas as pd
 import ribasim
 
+from ribasim_testmodels.utils import offset_spatial_inplace
+
 
 def user_model():
     """Create a user test model with static and dynamic users on the same basin."""
@@ -1165,9 +1167,17 @@ def main_network_with_subnetworks_model():
         endtime="2021-01-01 00:00:00",
     )
 
-    model.merge(subnetwork_model(), offset_spatial=(0.0, 3.0))
-    model.merge(fractional_flow_subnetwork_model(), offset_spatial=(14.0, 3.0))
-    model.merge(looped_subnetwork_model(), offset_spatial=(26.0, 3.0))
+    subnetwork_1 = subnetwork_model()
+    subnetwork_2 = fractional_flow_subnetwork_model()
+    subnetwork_3 = looped_subnetwork_model()
+
+    offset_spatial_inplace(subnetwork_1, (0.0, 3.0))
+    offset_spatial_inplace(subnetwork_2, (14.0, 3.0))
+    offset_spatial_inplace(subnetwork_3, (26.0, 3.0))
+
+    model.smart_merge(subnetwork_1)
+    model.smart_merge(subnetwork_2)
+    model.smart_merge(subnetwork_3)
 
     # Connection edges
     from_id = np.array([2, 6, 10])
diff --git a/python/ribasim_testmodels/ribasim_testmodels/utils.py b/python/ribasim_testmodels/ribasim_testmodels/utils.py
new file mode 100644
index 000000000..326670295
--- /dev/null
+++ b/python/ribasim_testmodels/ribasim_testmodels/utils.py
@@ -0,0 +1,6 @@
+def offset_spatial_inplace(model, offset: tuple[float, float]):
+    """Translate the geometry of a model with the given offset."""
+    network = model.network
+    network.edge.df.geometry = network.edge.df.geometry.translate(*offset)
+    network.node.df.geometry = network.node.df.geometry.translate(*offset)
+    return

From 31a37c10b67fcdfadb808148fe52645f143bb59b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Jan 2024 14:23:13 +0100
Subject: [PATCH 20/20] Add offset to to indices of new edges

---
 python/ribasim/ribasim/model.py | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index a75b7cb5a..7da70ff7a 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -61,6 +61,12 @@ def n_nodes(self):
 
         return n
 
+    def max_node_id(self) -> int:
+        return self.node.df.index.max()
+
+    def max_edge_id(self) -> int:
+        return self.edge.df.index.max()
+
     def offset_allocation_network_ids(
         self, offset_allocation_network_id: int, inplace: bool = True
     ) -> "Network":
@@ -86,13 +92,17 @@ def add_edges(
         else:
             network = deepcopy(self)
 
+        offset_edge_id = self.max_edge_id() + 1
+
         df = pd.DataFrame(
             data={
                 "from_node_id": from_node_id,
                 "to_node_id": to_node_id,
                 "edge_type": edge_type,
-            }
+            },
+            index=np.arange(offset_edge_id, offset_edge_id + len(edge_type)),
         )
+        print(df)
 
         df["geometry"] = df.apply(
             (
@@ -435,12 +445,6 @@ def reset_contextvar(self) -> "Model":
         context_file_loading.set({})
         return self
 
-    def max_node_id(self) -> int:
-        return self.network.node.df.index.max()
-
-    def max_edge_id(self) -> int:
-        return self.network.edge.df.index.max()
-
     def max_allocation_network_id(self) -> int:
         m = self.network.node.df.allocation_network_id.max()
         if pd.isna(m):
@@ -461,8 +465,8 @@ def smart_merge(
         nodes_model = self.nodes()
         nodes_added = model_added.nodes()
 
-        offset_node_id = self.max_node_id()
-        offset_edge_id = self.max_edge_id()
+        offset_node_id = self.network.max_node_id()
+        offset_edge_id = self.network.max_edge_id()
         offset_allocation_network_id = self.max_allocation_network_id()
 
         network_added = nodes_added["network"].offset_allocation_network_ids(