From 2a8de52d8b41d99980c73044315f2f3c1c988313 Mon Sep 17 00:00:00 2001 From: Maarten Pronk Date: Thu, 14 Mar 2024 17:39:32 +0100 Subject: [PATCH] Fix generation & mypy. --- core/src/schema.jl | 34 +++++++++++ python/ribasim/ribasim/schemas.py | 61 +++++++++---------- .../ribasim_testmodels/basic.py | 3 +- 3 files changed, 66 insertions(+), 32 deletions(-) diff --git a/core/src/schema.jl b/core/src/schema.jl index d7c1722d8..c7386b42f 100644 --- a/core/src/schema.jl +++ b/core/src/schema.jl @@ -7,12 +7,16 @@ @schema "ribasim.basin.profile" BasinProfile @schema "ribasim.basin.state" BasinState @schema "ribasim.basin.subgrid" BasinSubgrid +@schema "ribasim.basin.boundaryconcentration" BasinBoundaryconcentration +@schema "ribasim.basin.concentration" BasinConcentration @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 @@ -96,6 +100,22 @@ end urban_runoff::Union{Missing, Float64} end +@version BasinBoundaryconcentrationV1 begin + node_id::Int32 + time::DateTime + substance::String + drainage_concentration::Union{Missing, Float64} + precipitation_concentration::Union{Missing, Float64} + urban_runoff_concentration::Union{Missing, Float64} +end + +@version BasinConcentrationV1 begin + node_id::Int32 + time::DateTime + substance::String + concentration::Float64 +end + @version BasinProfileV1 begin node_id::Int32 area::Float64 @@ -132,6 +152,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} @@ -144,6 +171,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} diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py index ad09d5dc1..cac1bab75 100644 --- a/python/ribasim/ribasim/schemas.py +++ b/python/ribasim/ribasim/schemas.py @@ -11,6 +11,22 @@ class Config: coerce = True +class BasinBoundaryconcentrationSchema(_BaseSchema): + node_id: Series[int] = pa.Field(nullable=False, default=0) + time: Series[Timestamp] = pa.Field(nullable=False) + substance: Series[str] = pa.Field(nullable=False) + drainage_concentration: Series[float] = pa.Field(nullable=True) + precipitation_concentration: Series[float] = pa.Field(nullable=True) + urban_runoff_concentration: Series[float] = pa.Field(nullable=True) + + +class BasinConcentrationSchema(_BaseSchema): + node_id: Series[int] = 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 BasinProfileSchema(_BaseSchema): node_id: Series[int] = pa.Field(nullable=False, default=0) area: Series[float] = pa.Field(nullable=False) @@ -48,23 +64,6 @@ class BasinTimeSchema(_BaseSchema): urban_runoff: Series[float] = pa.Field(nullable=True) -# No capital for concentration to get Basin / boundaryconcentration -class BasinBoundaryconcentrationSchema(_BaseSchema): - node_id: Series[int] = pa.Field(nullable=False, default=0) - time: Series[Timestamp] = pa.Field(nullable=False) - substance: Series[str] = pa.Field(nullable=False) - drainage_concentration: Series[float] = pa.Field(nullable=True) - precipitation_concentration: Series[float] = pa.Field(nullable=True) - urban_runoff_concentration: Series[float] = pa.Field(nullable=True) - - -class BasinConcentrationSchema(_BaseSchema): - node_id: Series[int] = 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 DiscreteControlConditionSchema(_BaseSchema): node_id: Series[int] = pa.Field(nullable=False, default=0) listen_node_type: Series[str] = pa.Field(nullable=False) @@ -80,6 +79,13 @@ class DiscreteControlLogicSchema(_BaseSchema): control_state: Series[str] = pa.Field(nullable=False) +class FlowBoundaryConcentrationSchema(_BaseSchema): + node_id: Series[int] = 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[int] = pa.Field(nullable=False, default=0) active: Series[pa.BOOL] = pa.Field(nullable=True) @@ -92,19 +98,19 @@ class FlowBoundaryTimeSchema(_BaseSchema): flow_rate: Series[float] = pa.Field(nullable=False) -class FlowBoundaryConcentrationSchema(_BaseSchema): - node_id: Series[int] = 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 FractionalFlowStaticSchema(_BaseSchema): node_id: Series[int] = pa.Field(nullable=False, default=0) fraction: Series[float] = pa.Field(nullable=False) control_state: Series[str] = pa.Field(nullable=True) +class LevelBoundaryConcentrationSchema(_BaseSchema): + node_id: Series[int] = 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[int] = pa.Field(nullable=False, default=0) active: Series[pa.BOOL] = pa.Field(nullable=True) @@ -117,13 +123,6 @@ class LevelBoundaryTimeSchema(_BaseSchema): level: Series[float] = pa.Field(nullable=False) -class LevelBoundaryConcentrationSchema(_BaseSchema): - node_id: Series[int] = 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 LevelDemandStaticSchema(_BaseSchema): node_id: Series[int] = pa.Field(nullable=False, default=0) min_level: Series[float] = pa.Field(nullable=False) diff --git a/python/ribasim_testmodels/ribasim_testmodels/basic.py b/python/ribasim_testmodels/ribasim_testmodels/basic.py index 87e6cf7ba..0c3c5eb7b 100644 --- a/python/ribasim_testmodels/ribasim_testmodels/basic.py +++ b/python/ribasim_testmodels/ribasim_testmodels/basic.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from pathlib import Path from typing import Any @@ -116,7 +117,7 @@ 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_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]