Merging models

core/src/bmi.jl

@@ -73,15 +73,9 @@
     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

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

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

python/ribasim/ribasim/geometry/edge.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from typing import Any
 
 import matplotlib.pyplot as plt
@@ -39,6 +40,18 @@
         Table describing the flow connections.
     """
 
+    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:
+            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()
 

python/ribasim/ribasim/geometry/node.py

@@ -1,11 +1,14 @@
 from collections.abc import Sequence
+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
@@ -59,6 +62,18 @@
 
         return node_id, node_type
 
+    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:
+            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]:
@@ -130,6 +145,47 @@
         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")
+
+        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"
+        ):
+            if allocation_subnetwork_id is None:
+                continue
+            elif allocation_subnetwork_id == 1:
+                contains_main_network = True
+                color = COLOR_MAIN_NETWORK
+            else:
+                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:
         """
         Plot the nodes. Each node type is given a separate marker.

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 (
@@ -12,6 +13,7 @@
 )
 
 import geopandas as gpd
+import numpy as np
 import pandas as pd
 import pandera as pa
 from pandera.typing import DataFrame
@@ -159,6 +161,19 @@
 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]):
@@ -202,6 +217,81 @@
 
         return node_ids
 
+    def offset_node_ids(
+        self, offset_node_id: int, inplace: bool = True
+    ) -> "TableModel[TableT]":
+        """Add the same offset to all node IDs."""
+        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 [
+                "node_id",
+                "from_node_id",
+                "to_node_id",
+                "listen_feature_id",
+            ]:
+                if hasattr(df, name_column):
+                    df[name_column] += offset_node_id
+
+        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
+    ) -> "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."
+
+        if inplace:
+            table = self
+        else:
+            table = deepcopy(self)
+
+        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
+
     @classmethod
     def _load(cls, filepath: Path | None) -> dict[str, Any]:
         db = context_file_loading.get().get("database")
@@ -430,6 +520,61 @@
         ids = self.node_ids()
         return list(ids), len(ids) * [self.get_input_type()]
 
+    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."""
+        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,
+                    field,
+                    table.offset_node_ids(offset_node_id),
+                )
+        return node
+
+    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."
+
+        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)
+                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":
+        """Delete all rows of the underlying tables whose node ID is in the given list."""
+        if inplace:
+            node = self
+        else:
+            node = deepcopy(self)
+
+        for field in node.fields():
+            attr = getattr(node, field)
+            if isinstance(attr, TableModel):
+                df = attr.df
+                if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)):
+                    df = df[~df.node_id.isin(node_ids)]  # type: ignore
+                    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(