Add waterquality (substance/concentration) tables (#1267)

Fixes #1125. Split from #1137

core/src/schema.jl

@@ -8,12 +8,17 @@
 @schema "ribasim.basin.profile" BasinProfile
 @schema "ribasim.basin.state" BasinState
 @schema "ribasim.basin.subgrid" BasinSubgrid
+@schema "ribasim.basin.concentration" BasinConcentration
+@schema "ribasim.basin.concentrationexternal" BasinConcentrationExternal
+@schema "ribasim.basin.concentrationstate" BasinConcentrationState
 @schema "ribasim.terminal.static" TerminalStatic
 @schema "ribasim.fractionalflow.static" FractionalFlowStatic
 @schema "ribasim.flowboundary.static" FlowBoundaryStatic
 @schema "ribasim.flowboundary.time" FlowBoundaryTime
+@schema "ribasim.flowboundary.concentration" FlowBoundaryConcentration
 @schema "ribasim.levelboundary.static" LevelBoundaryStatic
 @schema "ribasim.levelboundary.time" LevelBoundaryTime
+@schema "ribasim.levelboundary.concentration" LevelBoundaryConcentration
 @schema "ribasim.linearresistance.static" LinearResistanceStatic
 @schema "ribasim.manningresistance.static" ManningResistanceStatic
 @schema "ribasim.pidcontrol.static" PidControlStatic
@@ -99,6 +104,21 @@ end
     urban_runoff::Union{Missing, Float64}
 end
 
+@version BasinConcentrationV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    drainage::Union{Missing, Float64}
+    precipitation::Union{Missing, Float64}
+end
+
+@version BasinConcentrationExternalV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    concentration::Union{Missing, Float64}
+end
+
 @version BasinProfileV1 begin
     node_id::Int32
     area::Float64
@@ -110,6 +130,12 @@ end
     level::Float64
 end
 
+@version BasinConcentrationStateV1 begin
+    node_id::Int32
+    substance::String
+    concentration::Union{Missing, Float64}
+end
+
 @version BasinSubgridV1 begin
     subgrid_id::Int32
     node_id::Int32
@@ -135,6 +161,13 @@ end
     level::Float64
 end
 
+@version LevelBoundaryConcentrationV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    concentration::Float64
+end
+
 @version FlowBoundaryStaticV1 begin
     node_id::Int32
     active::Union{Missing, Bool}
@@ -147,6 +180,13 @@ end
     flow_rate::Float64
 end
 
+@version FlowBoundaryConcentrationV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    concentration::Float64
+end
+
 @version LinearResistanceStaticV1 begin
     node_id::Int32
     active::Union{Missing, Bool}

python/ribasim/ribasim/config.py

@@ -16,6 +16,9 @@
 
 # These schemas are autogenerated
 from ribasim.schemas import (
+    BasinConcentrationExternalSchema,
+    BasinConcentrationSchema,
+    BasinConcentrationStateSchema,
     BasinProfileSchema,
     BasinStateSchema,
     BasinStaticSchema,
@@ -24,11 +27,13 @@
     DiscreteControlConditionSchema,
     DiscreteControlLogicSchema,
     DiscreteControlVariableSchema,
+    FlowBoundaryConcentrationSchema,
     FlowBoundaryStaticSchema,
     FlowBoundaryTimeSchema,
     FlowDemandStaticSchema,
     FlowDemandTimeSchema,
     FractionalFlowStaticSchema,
+    LevelBoundaryConcentrationSchema,
     LevelBoundaryStaticSchema,
     LevelBoundaryTimeSchema,
     LevelDemandStaticSchema,
@@ -189,6 +194,10 @@ class LevelBoundary(MultiNodeModel):
         default_factory=TableModel[LevelBoundaryTimeSchema],
         json_schema_extra={"sort_keys": ["node_id", "time"]},
     )
+    concentration: TableModel[LevelBoundaryConcentrationSchema] = Field(
+        default_factory=TableModel[LevelBoundaryConcentrationSchema],
+        json_schema_extra={"sort_keys": ["node_id", "substance", "time"]},
+    )
 
 
 class Pump(MultiNodeModel):
@@ -240,6 +249,10 @@ class FlowBoundary(MultiNodeModel):
         default_factory=TableModel[FlowBoundaryTimeSchema],
         json_schema_extra={"sort_keys": ["node_id", "time"]},
     )
+    concentration: TableModel[FlowBoundaryConcentrationSchema] = Field(
+        default_factory=TableModel[FlowBoundaryConcentrationSchema],
+        json_schema_extra={"sort_keys": ["node_id", "substance", "time"]},
+    )
 
 
 class FlowDemand(MultiNodeModel):
@@ -278,6 +291,18 @@ class Basin(MultiNodeModel):
         default_factory=SpatialTableModel[BasinAreaSchema],
         json_schema_extra={"sort_keys": ["node_id"]},
     )
+    concentration: TableModel[BasinConcentrationSchema] = Field(
+        default_factory=TableModel[BasinConcentrationSchema],
+        json_schema_extra={"sort_keys": ["node_id", "substance", "time"]},
+    )
+    concentration_external: TableModel[BasinConcentrationExternalSchema] = Field(
+        default_factory=TableModel[BasinConcentrationExternalSchema],
+        json_schema_extra={"sort_keys": ["node_id", "substance", "time"]},
+    )
+    concentration_state: TableModel[BasinConcentrationStateSchema] = Field(
+        default_factory=TableModel[BasinConcentrationStateSchema],
+        json_schema_extra={"sort_keys": ["node_id", "substance"]},
+    )
 
 
 class ManningResistance(MultiNodeModel):

python/ribasim/ribasim/nodes/basin.py

@@ -4,14 +4,25 @@
 from ribasim.geometry.area import BasinAreaSchema
 from ribasim.input_base import TableModel
 from ribasim.schemas import (
+    BasinConcentrationExternalSchema,
+    BasinConcentrationSchema,
+    BasinConcentrationStateSchema,
     BasinProfileSchema,
     BasinStateSchema,
     BasinStaticSchema,
     BasinSubgridSchema,
     BasinTimeSchema,
 )
 
-__all__ = ["Static", "Time", "State", "Profile", "Subgrid", "Area"]
+__all__ = [
+    "Static",
+    "Time",
+    "State",
+    "Profile",
+    "Subgrid",
+    "Area",
+    "Concentration",
+]
 
 
 class Static(TableModel[BasinStaticSchema]):
@@ -42,3 +53,18 @@ def __init__(self, **kwargs):
 class Area(TableModel[BasinAreaSchema]):
     def __init__(self, **kwargs):
         super().__init__(df=GeoDataFrame(dict(**kwargs)))
+
+
+class Concentration(TableModel[BasinConcentrationSchema]):
+    def __init__(self, **kwargs):
+        super().__init__(df=GeoDataFrame(dict(**kwargs)))
+
+
+class ConcentrationExternal(TableModel[BasinConcentrationExternalSchema]):
+    def __init__(self, **kwargs):
+        super().__init__(df=GeoDataFrame(dict(**kwargs)))
+
+
+class ConcentrationState(TableModel[BasinConcentrationStateSchema]):
+    def __init__(self, **kwargs):
+        super().__init__(df=GeoDataFrame(dict(**kwargs)))

python/ribasim/ribasim/nodes/flow_boundary.py

@@ -2,11 +2,12 @@
 
 from ribasim.input_base import TableModel
 from ribasim.schemas import (
+    FlowBoundaryConcentrationSchema,
     FlowBoundaryStaticSchema,
     FlowBoundaryTimeSchema,
 )
 
-__all__ = ["Static", "Time"]
+__all__ = ["Static", "Time", "Concentration"]
 
 
 class Static(TableModel[FlowBoundaryStaticSchema]):
@@ -17,3 +18,8 @@ def __init__(self, **kwargs):
 class Time(TableModel[FlowBoundaryTimeSchema]):
     def __init__(self, **kwargs):
         super().__init__(df=DataFrame(dict(**kwargs)))
+
+
+class Concentration(TableModel[FlowBoundaryConcentrationSchema]):
+    def __init__(self, **kwargs):
+        super().__init__(df=DataFrame(dict(**kwargs)))

python/ribasim/ribasim/nodes/level_boundary.py

@@ -2,11 +2,12 @@
 
 from ribasim.input_base import TableModel
 from ribasim.schemas import (
+    LevelBoundaryConcentrationSchema,
     LevelBoundaryStaticSchema,
     LevelBoundaryTimeSchema,
 )
 
-__all__ = ["Static", "Time"]
+__all__ = ["Static", "Time", "Concentration"]
 
 
 class Static(TableModel[LevelBoundaryStaticSchema]):
@@ -17,3 +18,8 @@ def __init__(self, **kwargs):
 class Time(TableModel[LevelBoundaryTimeSchema]):
     def __init__(self, **kwargs):
         super().__init__(df=DataFrame(dict(**kwargs)))
+
+
+class Concentration(TableModel[LevelBoundaryConcentrationSchema]):
+    def __init__(self, **kwargs):
+        super().__init__(df=DataFrame(dict(**kwargs)))

python/ribasim/ribasim/schemas.py

@@ -11,6 +11,27 @@ class Config:
         coerce = True
 
 
+class BasinConcentrationExternalSchema(_BaseSchema):
+    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    concentration: Series[float] = pa.Field(nullable=True)
+
+
+class BasinConcentrationStateSchema(_BaseSchema):
+    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    substance: Series[str] = pa.Field(nullable=False)
+    concentration: Series[float] = pa.Field(nullable=True)
+
+
+class BasinConcentrationSchema(_BaseSchema):
+    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    drainage: Series[float] = pa.Field(nullable=True)
+    precipitation: Series[float] = pa.Field(nullable=True)
+
+
 class BasinProfileSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     area: Series[float] = pa.Field(nullable=False)
@@ -70,6 +91,13 @@ class DiscreteControlVariableSchema(_BaseSchema):
     look_ahead: Series[float] = pa.Field(nullable=True)
 
 
+class FlowBoundaryConcentrationSchema(_BaseSchema):
+    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    concentration: Series[float] = pa.Field(nullable=False)
+
+
 class FlowBoundaryStaticSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     active: Series[pa.BOOL] = pa.Field(nullable=True)
@@ -101,6 +129,13 @@ class FractionalFlowStaticSchema(_BaseSchema):
     control_state: Series[str] = pa.Field(nullable=True)
 
 
+class LevelBoundaryConcentrationSchema(_BaseSchema):
+    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    concentration: Series[float] = pa.Field(nullable=False)
+
+
 class LevelBoundaryStaticSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     active: Series[pa.BOOL] = pa.Field(nullable=True)

python/ribasim_testmodels/ribasim_testmodels/basic.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from pathlib import Path
 from typing import Any
 
@@ -37,6 +38,16 @@ def basic_model() -> ribasim.Model:
             potential_evaporation=[0.001 / 86400], precipitation=[0.002 / 86400]
         ),
         basin.State(level=[0.04471158417652035]),
+        basin.Concentration(
+            time="2020-01-01 00:00:00",
+            substance=["Cl"],
+            drainage=[0.0],
+            precipitation=[0.0],
+        ),
+        basin.ConcentrationState(substance=["Cl"], concentration=[0.0]),
+        basin.ConcentrationExternal(
+            time="2020-01-01 00:00:00", substance=["Cl"], concentration=[0.0]
+        ),
     ]
     node_ids = [1, 3, 6, 9]
     node_geometries = [
@@ -107,16 +118,33 @@ def basic_model() -> ribasim.Model:
     model.pump.add(Node(7, Point(4.0, 1.0)), [pump.Static(flow_rate=[0.5 / 3600])])
 
     # Setup flow boundary
-    flow_boundary_data = [flow_boundary.Static(flow_rate=[1e-4])]
+    flow_boundary_data: Sequence[TableModel[Any]] = [
+        flow_boundary.Static(flow_rate=[1e-4]),
+        flow_boundary.Concentration(
+            time="2020-01-01 00:00:00", substance=["Tracer"], concentration=[1.0]
+        ),
+    ]
     model.flow_boundary.add(Node(15, Point(3.0, 3.0)), flow_boundary_data)
     model.flow_boundary.add(Node(16, Point(0.0, 1.0)), flow_boundary_data)
 
     # Setup level boundary
     model.level_boundary.add(
-        Node(11, Point(2.0, 2.0)), [level_boundary.Static(level=[1.0])]
+        Node(11, Point(2.0, 2.0)),
+        [
+            level_boundary.Static(level=[1.0]),
+            level_boundary.Concentration(
+                time="2020-01-01 00:00:00", substance=["Cl"], concentration=[34.0]
+            ),
+        ],
     )
     model.level_boundary.add(
-        Node(17, Point(6.0, 1.0)), [level_boundary.Static(level=[1.5])]
+        Node(17, Point(6.0, 1.0)),
+        [
+            level_boundary.Static(level=[1.5]),
+            level_boundary.Concentration(
+                time="2020-01-01 00:00:00", substance=["Cl"], concentration=[34.0]
+            ),
+        ],
     )
 
     # Setup terminal
@@ -340,9 +368,9 @@ def outlet_model():
         [
             level_boundary.Time(
                 time=[
-                    "2020-01-01",
-                    "2020-06-01",
-                    "2021-01-01",
+                    "2020-01-01 00:00:00",
+                    "2020-06-01 00:00:00",
+                    "2021-01-01 00:00:00",
                 ],
                 level=[1.0, 3.0, 3.0],
             )