Python edge validation

python/ribasim/ribasim/geometry/edge.py

@@ -15,6 +15,12 @@
 
 from ribasim.input_base import SpatialTableModel
 from ribasim.utils import UsedIDs
+from ribasim.validation import (
+    can_connect,
+    connectivity,
+    control_edge_amount,
+    flow_edge_amount,
+)
 
 from .base import _GeoBaseSchema
 
@@ -82,6 +88,11 @@ def add(
             the allocation algorithm, and should thus not be set for every edge in a subnetwork.
         **kwargs : Dict
         """
+        if not can_connect(from_node.node_type, to_node.node_type):
+            raise ValueError(
+                f"Node of type {to_node.node_type} cannot be downstream of node of type {from_node.node_type}. Possible downstream node: {connectivity[from_node.node_type]}."
+            )
+
         geometry_to_append = (
             [LineString([from_node.geometry, to_node.geometry])]
             if geometry is None
@@ -112,13 +123,62 @@ def add(
             index=pd.Index([edge_id], name="edge_id"),
         )
 
+        self._validate_edge(to_node, from_node, edge_type)
+
         self.df = GeoDataFrame[EdgeSchema](pd.concat([self.df, table_to_append]))
         if self.df.duplicated(subset=["from_node_id", "to_node_id"]).any():
             raise ValueError(
                 f"Edges have to be unique, but edge with from_node_id {from_node.node_id} to_node_id {to_node.node_id} already exists."
             )
         self._used_edge_ids.add(edge_id)
 
+    def _validate_edge(self, to_node: NodeData, from_node: NodeData, edge_type: str):
+        assert self.df is not None
+        in_flow_neighbor: int = self.df.loc[
+            (self.df["to_node_id"] == to_node.node_id)
+            & (self.df["edge_type"] == "flow")
+        ].shape[0]
+
+        out_flow_neighbor: int = self.df.loc[
+            (self.df["from_node_id"] == from_node.node_id)
+            & (self.df["edge_type"] == "flow")
+        ].shape[0]
+        # validation on neighbor amount
+        if (in_flow_neighbor >= flow_edge_amount[to_node.node_type][1]) & (
+            edge_type == "flow"
+        ):
+            raise ValueError(
+                f"Node {to_node.node_id} can have at most {flow_edge_amount[to_node.node_type][1]} flow edge inneighbor(s) (got {in_flow_neighbor})"
+            )
+        if (out_flow_neighbor >= flow_edge_amount[from_node.node_type][3]) & (
+            edge_type == "flow"
+        ):
+            raise ValueError(
+                f"Node {from_node.node_id} can have at most {flow_edge_amount[from_node.node_type][3]} flow edge outneighbor(s) (got {out_flow_neighbor})"
+            )
+
+        in_control_neighbor: int = self.df.loc[
+            (self.df["to_node_id"] == to_node.node_id)
+            & (self.df["edge_type"] == "control")
+        ].shape[0]
+        out_control_neighbor: int = self.df.loc[
+            (self.df["from_node_id"] == from_node.node_id)
+            & (self.df["edge_type"] == "control")
+        ].shape[0]
+
+        if (in_control_neighbor >= control_edge_amount[to_node.node_type][1]) & (
+            edge_type == "control"
+        ):
+            raise ValueError(
+                f"Node {to_node.node_id} can have at most {control_edge_amount[to_node.node_type][1]} control edge inneighbor(s) (got {in_control_neighbor})"
+            )
+        if (out_control_neighbor >= control_edge_amount[from_node.node_type][3]) & (
+            edge_type == "control"
+        ):
+            raise ValueError(
+                f"Node {from_node.node_id} can have at most {control_edge_amount[from_node.node_type][3]} control edge outneighbor(s) (got {out_control_neighbor})"
+            )
+
     def _get_where_edge_type(self, edge_type: str) -> NDArray[np.bool_]:
         assert self.df is not None
         return (self.df.edge_type == edge_type).to_numpy()

python/ribasim/ribasim/model.py

@@ -1,4 +1,5 @@
 import datetime
+import logging
 from collections.abc import Generator
 from os import PathLike
 from pathlib import Path
@@ -58,6 +59,7 @@
     _node_lookup_numpy,
     _time_in_ns,
 )
+from ribasim.validation import control_edge_amount, flow_edge_amount
 
 try:
     import xugrid
@@ -100,6 +102,7 @@ class Model(FileModel):
     user_demand: UserDemand = Field(default_factory=UserDemand)
 
     edge: EdgeTable = Field(default_factory=EdgeTable)
+    validation: bool = Field(default=True)
 
     _used_node_ids: UsedIDs = PrivateAttr(default_factory=UsedIDs)
 
@@ -265,8 +268,10 @@ def write(self, filepath: str | PathLike[str]) -> Path:
         filepath : str | PathLike[str]
             A file path with .toml extension.
         """
-        # TODO
-        # self.validate_model()
+        # Skip validation if the model name starts with "invalid"
+        if self.validation:
+            self._validate_model()
+
         filepath = Path(filepath)
         self.filepath = filepath
         if not filepath.suffix == ".toml":
@@ -280,6 +285,107 @@ def write(self, filepath: str | PathLike[str]) -> Path:
         context_file_writing.set({})
         return fn
 
+    def _validate_model(self):
+        df_edge = self.edge.df
+        df_chunks = [node.node.df.set_crs(self.crs) for node in self._nodes()]
+        df_node = pd.concat(df_chunks)
+
+        df_graph = df_edge
+        # Join df_edge with df_node to get to_node_type
+        df_graph = df_graph.join(
+            df_node[["node_type"]], on="from_node_id", how="left", rsuffix="_from"
+        )
+        df_graph = df_graph.rename(columns={"node_type": "from_node_type"})
+
+        df_graph = df_graph.join(
+            df_node[["node_type"]], on="to_node_id", how="left", rsuffix="_to"
+        )
+        df_graph = df_graph.rename(columns={"node_type": "to_node_type"})
+
+        if not self._check_neighbors(
+            df_graph, flow_edge_amount, "flow", df_node["node_type"]
+        ):
+            raise ValueError("Minimum flow inneighbor or outneighbor unsatisfied")
+        if not self._check_neighbors(
+            df_graph, control_edge_amount, "control", df_node["node_type"]
+        ):
+            raise ValueError("Minimum control inneighbor or outneighbor unsatisfied")
+
+    def _check_neighbors(
+        self,
+        df_graph: pd.DataFrame,
+        edge_amount: dict[str, list[int]],
+        edge_type: str,
+        nodes,
+    ) -> bool:
+        is_valid = True
+        # Count edge neighbor
+        df_graph = df_graph.loc[df_graph["edge_type"] == edge_type]
+
+        # check from node's neighbor
+        from_node_count = (
+            df_graph.groupby("from_node_id").size().reset_index(name="from_node_count")  # type: ignore
+        )
+
+        df_result = (
+            df_graph[["from_node_id", "from_node_type"]]
+            .drop_duplicates()
+            .merge(from_node_count, on="from_node_id", how="left")
+        )
+
+        df_result = df_result[["from_node_id", "from_node_count", "from_node_type"]]
+        for index, node in enumerate(nodes):
+            if nodes.index[index] not in df_result["from_node_id"].to_numpy():
+                new_row = {
+                    "from_node_id": nodes.index[index],
+                    "from_node_count": 0,
+                    "from_node_type": node,
+                }
+                df_result = pd.concat(
+                    [df_result, pd.DataFrame([new_row])], ignore_index=True
+                )
+        for _, row in df_result.iterrows():
+            # from node's outneighbor
+            try:
+                if row["from_node_count"] < edge_amount[row["from_node_type"]][2]:
+                    is_valid = False
+                    raise ValueError(
+                        f"Node {row['from_node_id']} must have at least {edge_amount[row['from_node_type']][2]} outneighbor(s) (got {row['from_node_count']})"
+                    )
+            except ValueError as e:
+                logging.error(e)
+        # check to node's neighbor
+        to_node_count = (
+            df_graph.groupby("to_node_id").size().reset_index(name="to_node_count")  # type: ignore
+        )
+        df_result = (
+            df_graph[["to_node_id", "to_node_type"]]
+            .drop_duplicates()
+            .merge(to_node_count, on="to_node_id", how="left")
+        )
+        df_result = df_result[["to_node_id", "to_node_count", "to_node_type"]]
+        for index, node in enumerate(nodes):
+            if nodes.index[index] not in df_result["to_node_id"].to_numpy():
+                new_row = {
+                    "to_node_id": nodes.index[index],
+                    "to_node_count": 0,
+                    "to_node_type": node,
+                }
+                df_result = pd.concat(
+                    [df_result, pd.DataFrame([new_row])], ignore_index=True
+                )
+
+        for _, row in df_result.iterrows():
+            try:
+                if row["to_node_count"] < edge_amount[row["to_node_type"]][0]:
+                    is_valid = False
+                    raise ValueError(
+                        f"Node {row['to_node_id']} must have at least {edge_amount[row['to_node_type']][0]} inneighbor(s) (got {row['to_node_count']})"
+                    )
+            except ValueError as e:
+                logging.error(e)
+        return is_valid
+
     @classmethod
     def _load(cls, filepath: Path | None) -> dict[str, Any]:
         context_file_loading.set({})

python/ribasim/ribasim/validation.py

@@ -0,0 +1,83 @@
+# Table for connectivity
+# "Basin": ["LinearResistance"] means that the downstream of basin can be LinearResistance only
+connectivity: dict[str, list[str]] = {
+    "Basin": [
+        "LinearResistance",
+        "ManningResistance",
+        "TabulatedRatingCurve",
+        "Pump",
+        "Outlet",
+        "UserDemand",
+    ],
+    "LinearResistance": ["Basin", "LevelBoundary"],
+    "ManningResistance": ["Basin"],
+    "TabulatedRatingCurve": ["Basin", "Terminal", "LevelBoundary"],
+    "LevelBoundary": ["LinearResistance", "Pump", "Outlet", "TabulatedRatingCurve"],
+    "FlowBoundary": ["Basin", "Terminal", "LevelBoundary"],
+    "Pump": ["Basin", "Terminal", "LevelBoundary"],
+    "Outlet": ["Basin", "Terminal", "LevelBoundary"],
+    "Terminal": [],
+    "DiscreteControl": [
+        "Pump",
+        "Outlet",
+        "TabulatedRatingCurve",
+        "LinearResistance",
+        "ManningResistance",
+        "PidControl",
+    ],
+    "ContinuousControl": ["Pump", "Outlet"],
+    "PidControl": ["Pump", "Outlet"],
+    "UserDemand": ["Basin", "Terminal", "LevelBoundary"],
+    "LevelDemand": ["Basin"],
+    "FlowDemand": [
+        "LinearResistance",
+        "ManningResistance",
+        "TabulatedRatingCurve",
+        "Pump",
+        "Outlet",
+    ],
+}
+
+
+# Function to validate connection
+def can_connect(node_up: str, node_down: str) -> bool:
+    if node_up in connectivity:
+        return node_down in connectivity[node_up]
+    return False
+
+
+flow_edge_amount: dict[str, list[int]] = {
+    "Basin": [0, int(1e9), 0, int(1e9)],
+    "LinearResistance": [1, 1, 1, 1],
+    "ManningResistance": [1, 1, 1, 1],
+    "TabulatedRatingCurve": [1, 1, 1, int(1e9)],
+    "LevelBoundary": [0, int(1e9), 0, int(1e9)],
+    "FlowBoundary": [0, 0, 1, int(1e9)],
+    "Pump": [1, 1, 1, int(1e9)],
+    "Outlet": [1, 1, 1, 1],
+    "Terminal": [1, int(1e9), 0, 0],
+    "DiscreteControl": [0, 0, 0, 0],
+    "ContinuousControl": [0, 0, 0, 0],
+    "PidControl": [0, 0, 0, 0],
+    "UserDemand": [1, 1, 1, 1],
+    "LevelDemand": [0, 0, 0, 0],
+    "FlowDemand": [0, 0, 0, 0],
+}
+
+control_edge_amount: dict[str, list[int]] = {
+    "Basin": [0, 1, 0, 0],
+    "LinearResistance": [0, 1, 0, 0],
+    "ManningResistance": [0, 1, 0, 0],
+    "TabulatedRatingCurve": [0, 1, 0, 0],
+    "LevelBoundary": [0, 0, 0, 0],
+    "FlowBoundary": [0, 0, 0, 0],
+    "Pump": [0, 1, 0, 0],
+    "Outlet": [0, 1, 0, 0],
+    "Terminal": [0, 0, 0, 0],
+    "DiscreteControl": [0, 0, 1, int(1e9)],
+    "ContinuousControl": [0, 0, 1, int(1e9)],
+    "PidControl": [0, 1, 1, 1],
+    "UserDemand": [0, 0, 0, 0],
+    "LevelDemand": [0, 0, 1, int(1e9)],
+    "FlowDemand": [0, 0, 1, 1],
+}

python/ribasim/tests/conftest.py

@@ -24,11 +24,21 @@ def bucket() -> ribasim.Model:
     return ribasim_testmodels.bucket_model()
 
 
+@pytest.fixture()
+def pid_control_equation() -> ribasim.Model:
+    return ribasim_testmodels.pid_control_equation_model()
+
+
 @pytest.fixture()
 def tabulated_rating_curve() -> ribasim.Model:
     return ribasim_testmodels.tabulated_rating_curve_model()
 
 
+@pytest.fixture()
+def outlet() -> ribasim.Model:
+    return ribasim_testmodels.outlet_model()
+
+
 @pytest.fixture()
 def backwater() -> ribasim.Model:
     return ribasim_testmodels.backwater_model()

python/ribasim/tests/test_io.py

@@ -130,7 +130,15 @@ def test_extra_spatial_columns():
             [basin.Profile(area=1000.0, level=[0.0, 1.0]), basin.State(level=[1.0])],
         )
     with pytest.raises(ValidationError):
-        model.edge.add(model.basin[1], model.user_demand[2], foo=1)
+        model.user_demand.add(
+            Node(4, Point(1, -0.5), meta_id=3),
+            [
+                user_demand.Static(
+                    demand=[1e-4], return_factor=0.9, min_level=0.9, priority=1
+                )
+            ],
+        )
+        model.edge.add(model.basin[1], model.user_demand[4], foo=1)
 
 
 def test_edge_autoincrement(basic):