Allow missing 'Basin / time' data (#1028)

Fixes #1027 Allows for missing values in basin.time data. --------- Co-authored-by: Martijn Visser <[email protected]>
Deltares · Feb 2, 2024 · 2dc5d74 · 2dc5d74
1 parent 793fdd7
commit 2dc5d74
Show file tree

Hide file tree

Showing 8 changed files with 196 additions and 54 deletions.
diff --git a/core/src/create.jl b/core/src/create.jl
@@ -485,10 +485,10 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
         current_area = DiffCache(current_area, chunk_sizes)
     end
 
-    precipitation = fill(NaN, length(node_id))
-    potential_evaporation = fill(NaN, length(node_id))
-    drainage = fill(NaN, length(node_id))
-    infiltration = fill(NaN, length(node_id))
+    precipitation = zeros(length(node_id))
+    potential_evaporation = zeros(length(node_id))
+    drainage = zeros(length(node_id))
+    infiltration = zeros(length(node_id))
     table = (; precipitation, potential_evaporation, drainage, infiltration)
 
     area, level, storage = create_storage_tables(db, config)

diff --git a/core/src/utils.jl b/core/src/utils.jl
@@ -625,12 +625,12 @@ end
 Update `table` at row index `i`, with the values of a given row.
 `table` must be a NamedTuple of vectors with all variables that must be loaded.
 The row must contain all the column names that are present in the table.
-If a value is NaN, it is not set.
+If a value is missing, it is not set.
 """
 function set_table_row!(table::NamedTuple, row, i::Int)::NamedTuple
     for (symbol, vector) in pairs(table)
         val = getproperty(row, symbol)
-        if !isnan(val)
+        if !ismissing(val)
             vector[i] = val
         end
     end
@@ -672,7 +672,7 @@ function set_current_value!(
     for (i, id) in enumerate(node_id)
         for (symbol, vector) in pairs(table)
             idx = findlast(
-                row -> row.node_id == id && !isnan(getproperty(row, symbol)),
+                row -> row.node_id == id && !ismissing(getproperty(row, symbol)),
                 pre_table,
             )
             if idx !== nothing

diff --git a/core/src/validation.jl b/core/src/validation.jl
@@ -185,11 +185,11 @@ end
 @version BasinTimeV1 begin
     node_id::Int
     time::DateTime
-    drainage::Float64
-    potential_evaporation::Float64
-    infiltration::Float64
-    precipitation::Float64
-    urban_runoff::Float64
+    drainage::Union{Missing, Float64}
+    potential_evaporation::Union{Missing, Float64}
+    infiltration::Union{Missing, Float64}
+    precipitation::Union{Missing, Float64}
+    urban_runoff::Union{Missing, Float64}
 end
 
 @version BasinProfileV1 begin

diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
@@ -98,6 +98,42 @@ end
     @test successful_retcode(model)
 end
 
+@testitem "leaky bucket model" begin
+    using SciMLBase: successful_retcode
+    import BasicModelInterface as BMI
+
+    toml_path = normpath(@__DIR__, "../../generated_testmodels/leaky_bucket/ribasim.toml")
+    @test ispath(toml_path)
+    model = Ribasim.Model(toml_path)
+    @test model isa Ribasim.Model
+
+    stor = model.integrator.u.storage
+    prec = model.integrator.p.basin.precipitation
+    evap = model.integrator.p.basin.potential_evaporation
+    drng = model.integrator.p.basin.drainage
+    infl = model.integrator.p.basin.infiltration
+    # The dynamic data has missings, but these are not set.
+    @test prec == [0.0]
+    @test evap == [0.0]
+    @test drng == [0.003]
+    @test infl == [0.0]
+    init_stor = 1000.0
+    @test stor == [init_stor]
+    BMI.update_until(model, 1.5 * 86400)
+    @test prec == [0.0]
+    @test evap == [0.0]
+    @test drng == [0.003]
+    @test infl == [0.001]
+    stor ≈ Float32[init_stor + 86400 * (0.003 * 1.5 - 0.001 * 0.5)]
+    BMI.update_until(model, 2.5 * 86400)
+    @test prec == [0.00]
+    @test evap == [0.0]
+    @test drng == [0.001]
+    @test infl == [0.002]
+    stor ≈ Float32[init_stor + 86400 * (0.003 * 2.0 + 0.001 * 0.5 - 0.001 - 0.002 * 0.5)]
+    @test successful_retcode(Ribasim.solve!(model))
+end
+
 @testitem "basic model" begin
     using Logging: Debug, with_logger
     using LoggingExtras

diff --git a/docs/schema/BasinTime.schema.json b/docs/schema/BasinTime.schema.json
@@ -14,24 +14,59 @@
             "type": "string"
         },
         "drainage": {
-            "format": "double",
-            "type": "number"
+            "format": "default",
+            "anyOf": [
+                {
+                    "type": "null"
+                },
+                {
+                    "type": "number"
+                }
+            ]
         },
         "potential_evaporation": {
-            "format": "double",
-            "type": "number"
+            "format": "default",
+            "anyOf": [
+                {
+                    "type": "null"
+                },
+                {
+                    "type": "number"
+                }
+            ]
         },
         "infiltration": {
-            "format": "double",
-            "type": "number"
+            "format": "default",
+            "anyOf": [
+                {
+                    "type": "null"
+                },
+                {
+                    "type": "number"
+                }
+            ]
         },
         "precipitation": {
-            "format": "double",
-            "type": "number"
+            "format": "default",
+            "anyOf": [
+                {
+                    "type": "null"
+                },
+                {
+                    "type": "number"
+                }
+            ]
         },
         "urban_runoff": {
-            "format": "double",
-            "type": "number"
+            "format": "default",
+            "anyOf": [
+                {
+                    "type": "null"
+                },
+                {
+                    "type": "number"
+                }
+            ]
         },
         "remarks": {
             "description": "a hack for pandera",
@@ -42,11 +77,6 @@
     },
     "required": [
         "node_id",
-        "time",
-        "drainage",
-        "potential_evaporation",
-        "infiltration",
-        "precipitation",
-        "urban_runoff"
+        "time"
     ]
 }
diff --git a/python/ribasim/ribasim/models.py b/python/ribasim/ribasim/models.py
@@ -44,11 +44,11 @@ class BasinSubgrid(BaseModel):
 class BasinTime(BaseModel):
     node_id: int
     time: datetime
-    drainage: float
-    potential_evaporation: float
-    infiltration: float
-    precipitation: float
-    urban_runoff: float
+    drainage: float | None = None
+    potential_evaporation: float | None = None
+    infiltration: float | None = None
+    precipitation: float | None = None
+    urban_runoff: float | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 

diff --git a/python/ribasim_testmodels/ribasim_testmodels/__init__.py b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
@@ -21,7 +21,7 @@
     outlet_model,
     tabulated_rating_curve_model,
 )
-from ribasim_testmodels.bucket import bucket_model
+from ribasim_testmodels.bucket import bucket_model, leaky_bucket_model
 from ribasim_testmodels.discrete_control import (
     flow_condition_model,
     level_boundary_condition_model,
@@ -54,37 +54,38 @@
 __all__ = [
     "allocation_example_model",
     "backwater_model",
-    "basic_model",
     "basic_arrow_model",
+    "basic_model",
     "basic_transient_model",
     "bucket_model",
-    "pump_discrete_control_model",
-    "flow_condition_model",
-    "tabulated_rating_curve_model",
-    "trivial_model",
-    "linear_resistance_model",
-    "rating_curve_model",
-    "manning_resistance_model",
-    "pid_control_model",
-    "misc_nodes_model",
-    "tabulated_rating_curve_control_model",
+    "discrete_control_of_pid_control_model",
     "dutch_waterways_model",
-    "invalid_qh_model",
     "flow_boundary_time_model",
-    "pid_control_equation_model",
-    "invalid_fractional_flow_model",
+    "flow_condition_model",
+    "fractional_flow_subnetwork_model",
     "invalid_discrete_control_model",
-    "level_setpoint_with_minmax_model",
     "invalid_edge_types_model",
-    "discrete_control_of_pid_control_model",
+    "invalid_fractional_flow_model",
+    "invalid_qh_model",
+    "leaky_bucket_model",
     "level_boundary_condition_model",
+    "level_setpoint_with_minmax_model",
+    "linear_resistance_model",
+    "looped_subnetwork_model",
+    "manning_resistance_model",
+    "minimal_subnetwork_model",
+    "misc_nodes_model",
     "outlet_model",
-    "user_model",
+    "pid_control_equation_model",
+    "pid_control_model",
+    "pump_discrete_control_model",
+    "rating_curve_model",
     "subnetwork_model",
-    "minimal_subnetwork_model",
-    "fractional_flow_subnetwork_model",
-    "looped_subnetwork_model",
+    "tabulated_rating_curve_control_model",
+    "tabulated_rating_curve_model",
+    "trivial_model",
     "two_basin_model",
+    "user_model",
 ]
 
 # provide a mapping from model name to its constructor, so we can iterate over all models

diff --git a/python/ribasim_testmodels/ribasim_testmodels/bucket.py b/python/ribasim_testmodels/ribasim_testmodels/bucket.py
@@ -76,3 +76,78 @@ def bucket_model() -> ribasim.Model:
         endtime="2021-01-01 00:00:00",
     )
     return model
+
+
+def leaky_bucket_model() -> ribasim.Model:
+    """Bucket model with dynamic forcing with missings."""
+
+    # Set up the nodes:
+    xy = np.array(
+        [
+            (400.0, 200.0),  # Basin
+        ]
+    )
+    node_xy = gpd.points_from_xy(x=xy[:, 0], y=xy[:, 1])
+    node_type = ["Basin"]
+    # Make sure the feature id starts at 1: explicitly give an index.
+    node = ribasim.Node(
+        df=gpd.GeoDataFrame(
+            data={"type": node_type},
+            index=pd.Index(np.arange(len(xy)) + 1, name="fid"),
+            geometry=node_xy,
+            crs="EPSG:28992",
+        )
+    )
+
+    # Setup the dummy edges:
+    from_id = np.array([], dtype=np.int64)
+    to_id = np.array([], dtype=np.int64)
+    lines = node.geometry_from_connectivity(from_id.tolist(), to_id.tolist())
+    edge = ribasim.Edge(
+        df=gpd.GeoDataFrame(
+            data={
+                "from_node_id": from_id,
+                "to_node_id": to_id,
+                "edge_type": len(from_id) * ["flow"],
+            },
+            geometry=lines,
+            crs="EPSG:28992",
+        )
+    )
+
+    # Setup the basins:
+    profile = pd.DataFrame(
+        data={
+            "node_id": [1, 1],
+            "area": [1000.0, 1000.0],
+            "level": [0.0, 1.0],
+        }
+    )
+
+    state = pd.DataFrame(
+        data={
+            "node_id": [1],
+            "level": [1.0],
+        }
+    )
+
+    time = pd.DataFrame(
+        data={
+            "time": pd.date_range("2020-01-01", "2020-01-05"),
+            "node_id": 1,
+            "drainage": [0.003, np.nan, 0.001, 0.002, 0.0],
+            "potential_evaporation": np.nan,
+            "infiltration": [np.nan, 0.001, 0.002, 0.0, 0.0],
+            "precipitation": np.nan,
+            "urban_runoff": 0.0,
+        }
+    )
+    basin = ribasim.Basin(profile=profile, time=time, state=state)
+
+    model = ribasim.Model(
+        network=ribasim.Network(node=node, edge=edge),
+        basin=basin,
+        starttime="2020-01-01 00:00:00",
+        endtime="2020-01-05 00:00:00",
+    )
+    return model