Prefix extra columns with meta_ (#794)

Fixes #576 Required, but seemingly unrelated, fixes: - The unused, but required RootModel is generated automatically, but had the wrong casing in fieldnames, leading to errors since switching to | union (#793) notation. - Utils required a Node import, leading to circular import errors when using utils elsewhere. Moved the functions to the Node class where they belong. --------- Co-authored-by: Martijn Visser <[email protected]>
Deltares · Nov 16, 2023 · 548c6eb · 548c6eb
1 parent a688b3d
commit 548c6eb
Show file tree

Hide file tree

Showing 22 changed files with 238 additions and 194 deletions.
diff --git a/docs/_quarto.yml b/docs/_quarto.yml
@@ -91,8 +91,3 @@ quartodoc:
         - Terminal
         - DiscreteControl
         - PidControl
-    - title: Utility functions
-      desc: Collection of utility functions.
-      contents:
-        - utils.geometry_from_connectivity
-        - utils.connectivity_from_geometry
diff --git a/docs/gen_schema.jl b/docs/gen_schema.jl
@@ -67,7 +67,7 @@ function gen_root_schema(TT::Vector, prefix = prefix, name = "root")
         "type" => "object",
     )
     for T in TT
-        tname = strip_prefix(T)
+        tname = lowercase(strip_prefix(T))
         schema["properties"][tname] = OrderedDict("\$ref" => "$tname.schema.json")
     end
     open(normpath(@__DIR__, "schema", "$(name).schema.json"), "w") do io

diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
@@ -361,7 +361,7 @@
     "to_id = np.array(\n",
     "    [2, 3, 4, 5, 8, 6, 7, 9, 9, 10, 12, 3, 13, 14, 6, 1, 17], dtype=np.int64\n",
     ")\n",
-    "lines = ribasim.utils.geometry_from_connectivity(node, from_id, to_id)\n",
+    "lines = node.geometry_from_connectivity(from_id, to_id)\n",
     "edge = ribasim.Edge(\n",
     "    df=gpd.GeoDataFrame(\n",
     "        data={\n",
@@ -705,7 +705,7 @@
     "\n",
     "edge_type = 6 * [\"flow\"] + 2 * [\"control\"]\n",
     "\n",
-    "lines = ribasim.utils.geometry_from_connectivity(node, from_id, to_id)\n",
+    "lines = node.geometry_from_connectivity(from_id, to_id)\n",
     "edge = ribasim.Edge(\n",
     "    df=gpd.GeoDataFrame(\n",
     "        data={\"from_node_id\": from_id, \"to_node_id\": to_id, \"edge_type\": edge_type},\n",
@@ -1132,7 +1132,7 @@
     "from_id = np.array([1, 2, 3, 4, 6, 5, 7], dtype=np.int64)\n",
     "to_id = np.array([2, 3, 4, 6, 2, 3, 6], dtype=np.int64)\n",
     "\n",
-    "lines = ribasim.utils.geometry_from_connectivity(node, from_id, to_id)\n",
+    "lines = node.geometry_from_connectivity(from_id, to_id)\n",
     "edge = ribasim.Edge(\n",
     "    df=gpd.GeoDataFrame(\n",
     "        data={\n",

diff --git a/docs/schema/root.schema.json b/docs/schema/root.schema.json
@@ -1,77 +1,77 @@
 {
     "$schema": "https://json-schema.org/draft/2020-12/schema",
     "properties": {
-        "BasinProfile": {
-            "$ref": "BasinProfile.schema.json"
+        "basinprofile": {
+            "$ref": "basinprofile.schema.json"
         },
-        "BasinState": {
-            "$ref": "BasinState.schema.json"
+        "basinstate": {
+            "$ref": "basinstate.schema.json"
         },
-        "BasinStatic": {
-            "$ref": "BasinStatic.schema.json"
+        "basinstatic": {
+            "$ref": "basinstatic.schema.json"
         },
-        "BasinTime": {
-            "$ref": "BasinTime.schema.json"
+        "basintime": {
+            "$ref": "basintime.schema.json"
         },
-        "DiscreteControlCondition": {
-            "$ref": "DiscreteControlCondition.schema.json"
+        "discretecontrolcondition": {
+            "$ref": "discretecontrolcondition.schema.json"
         },
-        "DiscreteControlLogic": {
-            "$ref": "DiscreteControlLogic.schema.json"
+        "discretecontrollogic": {
+            "$ref": "discretecontrollogic.schema.json"
         },
-        "Edge": {
-            "$ref": "Edge.schema.json"
+        "edge": {
+            "$ref": "edge.schema.json"
         },
-        "FlowBoundaryStatic": {
-            "$ref": "FlowBoundaryStatic.schema.json"
+        "flowboundarystatic": {
+            "$ref": "flowboundarystatic.schema.json"
         },
-        "FlowBoundaryTime": {
-            "$ref": "FlowBoundaryTime.schema.json"
+        "flowboundarytime": {
+            "$ref": "flowboundarytime.schema.json"
         },
-        "FractionalFlowStatic": {
-            "$ref": "FractionalFlowStatic.schema.json"
+        "fractionalflowstatic": {
+            "$ref": "fractionalflowstatic.schema.json"
         },
-        "LevelBoundaryStatic": {
-            "$ref": "LevelBoundaryStatic.schema.json"
+        "levelboundarystatic": {
+            "$ref": "levelboundarystatic.schema.json"
         },
-        "LevelBoundaryTime": {
-            "$ref": "LevelBoundaryTime.schema.json"
+        "levelboundarytime": {
+            "$ref": "levelboundarytime.schema.json"
         },
-        "LinearResistanceStatic": {
-            "$ref": "LinearResistanceStatic.schema.json"
+        "linearresistancestatic": {
+            "$ref": "linearresistancestatic.schema.json"
         },
-        "ManningResistanceStatic": {
-            "$ref": "ManningResistanceStatic.schema.json"
+        "manningresistancestatic": {
+            "$ref": "manningresistancestatic.schema.json"
         },
-        "Node": {
-            "$ref": "Node.schema.json"
+        "node": {
+            "$ref": "node.schema.json"
         },
-        "OutletStatic": {
-            "$ref": "OutletStatic.schema.json"
+        "outletstatic": {
+            "$ref": "outletstatic.schema.json"
         },
-        "PidControlStatic": {
-            "$ref": "PidControlStatic.schema.json"
+        "pidcontrolstatic": {
+            "$ref": "pidcontrolstatic.schema.json"
         },
-        "PidControlTime": {
-            "$ref": "PidControlTime.schema.json"
+        "pidcontroltime": {
+            "$ref": "pidcontroltime.schema.json"
         },
-        "PumpStatic": {
-            "$ref": "PumpStatic.schema.json"
+        "pumpstatic": {
+            "$ref": "pumpstatic.schema.json"
         },
-        "TabulatedRatingCurveStatic": {
-            "$ref": "TabulatedRatingCurveStatic.schema.json"
+        "tabulatedratingcurvestatic": {
+            "$ref": "tabulatedratingcurvestatic.schema.json"
         },
-        "TabulatedRatingCurveTime": {
-            "$ref": "TabulatedRatingCurveTime.schema.json"
+        "tabulatedratingcurvetime": {
+            "$ref": "tabulatedratingcurvetime.schema.json"
         },
-        "TerminalStatic": {
-            "$ref": "TerminalStatic.schema.json"
+        "terminalstatic": {
+            "$ref": "terminalstatic.schema.json"
         },
-        "UserStatic": {
-            "$ref": "UserStatic.schema.json"
+        "userstatic": {
+            "$ref": "userstatic.schema.json"
         },
-        "UserTime": {
-            "$ref": "UserTime.schema.json"
+        "usertime": {
+            "$ref": "usertime.schema.json"
         }
     },
     "$id": "https://deltares.github.io/Ribasim/schema/root.schema.json",

diff --git a/pixi.toml b/pixi.toml
@@ -77,7 +77,7 @@ tests = { depends_on = ["lint", "test-ribasim-python", "test-ribasim-core"] }
 generate-schema = { cmd = "julia --project=docs docs/gen_schema.jl", depends_on = [
     "instantiate-julia",
 ] }
-generate-python = "datamodel-codegen --use-union-operator --use-title-as-name --use-double-quotes --disable-timestamp --use-default --strict-nullable --input-file-type=jsonschema --input docs/schema/root.schema.json --output python/ribasim/ribasim/models.py"
+generate-python = "datamodel-codegen --output-model-type pydantic_v2.BaseModel --base-class ribasim.input_base.BaseModel --use-union-operator --use-title-as-name --use-double-quotes --disable-timestamp --use-default --strict-nullable --input-file-type=jsonschema --input docs/schema/root.schema.json --output python/ribasim/ribasim/models.py"
 codegen = { depends_on = ["generate-schema", "generate-python", "lint"] }
 # Publish
 build-ribasim-python-wheel = { cmd = "rm --recursive --force dist && python -m build && twine check dist/*", cwd = "python/ribasim" }

diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
@@ -1,9 +1,12 @@
+from collections.abc import Sequence
 from typing import Any
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import pandera as pa
+import shapely
+from numpy.typing import NDArray
 from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
 
@@ -56,6 +59,77 @@ def node_ids_and_types(*nodes):
 
         return node_id, node_type
 
+    def geometry_from_connectivity(
+        self, from_id: Sequence[int], to_id: Sequence[int]
+    ) -> NDArray[Any]:
+        """
+        Create edge shapely geometries from connectivities.
+
+        Parameters
+        ----------
+        node : Ribasim.Node
+        from_id : Sequence[int]
+            First node of every edge.
+        to_id : Sequence[int]
+            Second node of every edge.
+
+        Returns
+        -------
+        edge_geometry : np.ndarray
+            Array of shapely LineStrings.
+        """
+        geometry = self.df["geometry"]
+        from_points = shapely.get_coordinates(geometry.loc[from_id])
+        to_points = shapely.get_coordinates(geometry.loc[to_id])
+        n = len(from_points)
+        vertices = np.empty((n * 2, 2), dtype=from_points.dtype)
+        vertices[0::2, :] = from_points
+        vertices[1::2, :] = to_points
+        indices = np.repeat(np.arange(n), 2)
+        return shapely.linestrings(coords=vertices, indices=indices)
+
+    def connectivity_from_geometry(
+        self, lines: NDArray[Any]
+    ) -> tuple[NDArray[Any], NDArray[Any]]:
+        """
+        Derive from_node_id and to_node_id for every edge in lines. LineStrings
+        may be used to connect multiple nodes in a sequence, but every linestring
+        vertex must also a node.
+
+        Parameters
+        ----------
+        node : Node
+        lines : np.ndarray
+            Array of shapely linestrings.
+
+        Returns
+        -------
+        from_node_id : np.ndarray of int
+        to_node_id : np.ndarray of int
+        """
+        node_index = self.df.index
+        node_xy = shapely.get_coordinates(self.df.geometry.values)
+        edge_xy = shapely.get_coordinates(lines)
+
+        xy = np.vstack([node_xy, edge_xy])
+        _, inverse = np.unique(xy, return_inverse=True, axis=0)
+        _, index, inverse = np.unique(
+            xy, return_index=True, return_inverse=True, axis=0
+        )
+        uniques_index = index[inverse]
+
+        node_node_id, edge_node_id = np.split(uniques_index, [len(node_xy)])
+        if not np.isin(edge_node_id, node_node_id).all():
+            raise ValueError(
+                "Edge lines contain coordinates that are not in the node layer. "
+                "Please ensure all edges are snapped to nodes exactly."
+            )
+
+        edge_node_id = edge_node_id.reshape((-1, 2))
+        from_id = node_index[edge_node_id[:, 0]].to_numpy()
+        to_id = node_index[edge_node_id[:, 1]].to_numpy()
+        return from_id, to_id
+
     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/input_base.py b/python/ribasim/ribasim/input_base.py
@@ -20,11 +20,13 @@
     ConfigDict,
     DirectoryPath,
     Field,
+    field_validator,
     model_serializer,
     model_validator,
 )
 
 from ribasim.types import FilePath
+from ribasim.utils import prefix_column
 
 __all__ = ("TableModel",)
 
@@ -146,6 +148,16 @@ 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)
 
+    @field_validator("df")
+    @classmethod
+    def prefix_extra_columns(cls, v: DataFrame[TableT]):
+        """Prefix extra columns with meta_."""
+        if isinstance(v, pd.DataFrame):
+            v.rename(
+                lambda x: prefix_column(x, cls.columns()), axis="columns", inplace=True
+            )
+        return v
+
     @model_serializer
     def set_model(self) -> Path | None:
         return self.filepath
@@ -256,14 +268,24 @@ def tableschema(cls) -> TableT:
         T: TableT = fieldtype.__args__[0]
         return T
 
-    def record(self):
+    @classmethod
+    def record(cls) -> type[PydanticBaseModel] | None:
         """Retrieve Pydantic Record used in Pandera Schema."""
-        T = self.tableschema()
-        return T.Config.dtype.type
+        T = cls.tableschema()
+        if hasattr(T.Config, "dtype"):
+            # We always set a PydanticBaseModel dtype (see schemas.py)
+            return T.Config.dtype.type  # type: ignore
+        else:
+            return None
 
-    def columns(self):
+    @classmethod
+    def columns(cls) -> list[str]:
         """Retrieve column names."""
-        return list(self.record().model_fields.keys())
+        T = cls.record()
+        if T is not None:
+            return list(T.model_fields.keys())
+        else:
+            return []
 
 
 class SpatialTableModel(TableModel[TableT], Generic[TableT]):

diff --git a/python/ribasim/ribasim/models.py b/python/ribasim/ribasim/models.py
@@ -5,7 +5,9 @@
 
 from datetime import datetime
 
-from pydantic import BaseModel, Field
+from pydantic import Field
+
+from ribasim.input_base import BaseModel
 
 
 class BasinProfile(BaseModel):
@@ -218,27 +220,27 @@ class UserTime(BaseModel):
 
 
 class Root(BaseModel):
-    BasinProfile: BasinProfile | None
-    BasinState: BasinState | None
-    BasinStatic: BasinStatic | None
-    BasinTime: BasinTime | None
-    DiscreteControlCondition: DiscreteControlCondition | None
-    DiscreteControlLogic: DiscreteControlLogic | None
-    Edge: Edge | None
-    FlowBoundaryStatic: FlowBoundaryStatic | None
-    FlowBoundaryTime: FlowBoundaryTime | None
-    FractionalFlowStatic: FractionalFlowStatic | None
-    LevelBoundaryStatic: LevelBoundaryStatic | None
-    LevelBoundaryTime: LevelBoundaryTime | None
-    LinearResistanceStatic: LinearResistanceStatic | None
-    ManningResistanceStatic: ManningResistanceStatic | None
-    Node: Node | None
-    OutletStatic: OutletStatic | None
-    PidControlStatic: PidControlStatic | None
-    PidControlTime: PidControlTime | None
-    PumpStatic: PumpStatic | None
-    TabulatedRatingCurveStatic: TabulatedRatingCurveStatic | None
-    TabulatedRatingCurveTime: TabulatedRatingCurveTime | None
-    TerminalStatic: TerminalStatic | None
-    UserStatic: UserStatic | None
-    UserTime: UserTime | None
+    basinprofile: BasinProfile | None = None
+    basinstate: BasinState | None = None
+    basinstatic: BasinStatic | None = None
+    basintime: BasinTime | None = None
+    discretecontrolcondition: DiscreteControlCondition | None = None
+    discretecontrollogic: DiscreteControlLogic | None = None
+    edge: Edge | None = None
+    flowboundarystatic: FlowBoundaryStatic | None = None
+    flowboundarytime: FlowBoundaryTime | None = None
+    fractionalflowstatic: FractionalFlowStatic | None = None
+    levelboundarystatic: LevelBoundaryStatic | None = None
+    levelboundarytime: LevelBoundaryTime | None = None
+    linearresistancestatic: LinearResistanceStatic | None = None
+    manningresistancestatic: ManningResistanceStatic | None = None
+    node: Node | None = None
+    outletstatic: OutletStatic | None = None
+    pidcontrolstatic: PidControlStatic | None = None
+    pidcontroltime: PidControlTime | None = None
+    pumpstatic: PumpStatic | None = None
+    tabulatedratingcurvestatic: TabulatedRatingCurveStatic | None = None
+    tabulatedratingcurvetime: TabulatedRatingCurveTime | None = None
+    terminalstatic: TerminalStatic | None = None
+    userstatic: UserStatic | None = None
+    usertime: UserTime | None = None