DiscreteControl based on storage (#2573)

Fixes #2569

---------

Co-authored-by: Martijn Visser <[email protected]>

Author: SouthEndMusic

core/src/parameter.jl

@@ -786,7 +786,7 @@ const StateTimeDependentCache{T} = @NamedTuple{
     u_prev_call::Vector{T},
 } where {T}
 
-@enumx CacheType flow_rate_pump flow_rate_outlet basin_level
+@enumx CacheType flow_rate_pump flow_rate_outlet basin_level basin_storage
 
 """
 A cache for intermediate results in `water_balance!` which depend only on the time `t`. A second version of this
@@ -849,6 +849,8 @@ function get_cache_vector(
         state_time_dependent_cache.current_flow_rate_outlet
     elseif type == CacheType.basin_level
         state_time_dependent_cache.current_level
+    elseif type == CacheType.basin_storage
+        state_time_dependent_cache.current_storage
     else
         error("Invalid cache type $type passed.")
     end

core/src/util.jl

@@ -509,6 +509,8 @@ function get_cache_ref(
 
     ref = if node_id.type == NodeType.Basin && variable == "level"
         CacheRef(; type = CacheType.basin_level, node_id.idx)
+    elseif node_id.type == NodeType.Basin && variable == "storage"
+        CacheRef(; type = CacheType.basin_storage, node_id.idx)
     elseif variable == "flow_rate" && node_id.type != NodeType.FlowBoundary
         if listen
             if node_id.type ∉ conservative_nodetypes

core/test/control_test.jl

@@ -381,3 +381,14 @@ end
     @test control.truth_state[3] == "F"
     @test basin6.level[findfirst(>=(t2), basin6.time)] <= 0.9 + 1e-2
 end
+
+@testitem "Storage condition" begin
+    toml_path =
+        normpath(@__DIR__, "../../generated_testmodels/storage_condition/ribasim.toml")
+    @test ispath(toml_path)
+    model = Ribasim.run(toml_path)
+    @test success(model)
+
+    storage = Ribasim.get_storages_and_levels(model).storage[1, :]
+    @test all(storage .< 7500 + 6)
+end

docs/changelog.qmd

@@ -32,6 +32,7 @@ Lastly, Ribasim is now multithreaded!
 - Warn when reading a new model with an old Ribasim Python version. [#2410](https://github.com/Deltares/Ribasim/pull/2410)
 - Add instructions to install using Pixi. [#2415](https://github.com/Deltares/Ribasim/pull/2415)
 - Add guide on updating Ribasim models. [#2489](https://github.com/Deltares/Ribasim/pull/2489)
+- Support DiscreteControl listening to a Basin's storage. [#2573](https://github.com/Deltares/Ribasim/pull/2573)
 
 ### Changed
 - Allow two incoming control links for Pump and Outlet. [#2531](https://github.com/Deltares/Ribasim/pull/2531)

docs/reference/node/discrete-control.qmd

@@ -57,11 +57,12 @@ column               | type     | unit       | restriction
 node_id              | Int32    | -          |
 compound_variable_id | Int32    | -          |
 listen_node_id       | Int32    | -          | cannot be a Junction
-variable             | String   | -          | must be "level" or "flow_rate"
+variable             | String   | -          | must be "level", "storage" or "flow_rate"
 weight               | Float64  | -          | (optional, default 1.0)
 look_ahead           | Float64  | $\text{s}$ | Only on transient boundary conditions, non-negative (optional, default 0.0).
 
 These variables can be listened to:
+- The storage of a Basin
 - The level of a Basin
 - The level of a LevelBoundary (supports look ahead)
 - The flow rate through one of these node types: Pump, Outlet, TabulatedRatingCurve, LinearResistance, ManningResistance

python/ribasim_testmodels/ribasim_testmodels/__init__.py

@@ -54,6 +54,7 @@
     level_boundary_condition_model,
     level_range_model,
     pump_discrete_control_model,
+    storage_condition_model,
     tabulated_rating_curve_control_model,
     transient_condition_model,
 )
@@ -149,6 +150,7 @@
     "pump_discrete_control_model",
     "rating_curve_model",
     "rating_curve_between_basins_model",
+    "storage_condition_model",
     "subnetwork_model",
     "subnetworks_with_sources_model",
     "tabulated_rating_curve_control_model",

python/ribasim_testmodels/ribasim_testmodels/discrete_control.py

@@ -496,6 +496,55 @@ def level_range_model() -> Model:
     return model
 
 
+def storage_condition_model() -> Model:
+    """Create a model with a discrete control condition based on the storage of a Basin."""
+    model = Model(
+        starttime="2020-01-01",
+        endtime="2021-01-01",
+        crs="EPSG:28992",
+    )
+
+    fb = model.flow_boundary.add(
+        Node(1, Point(0, 0)), [flow_boundary.Static(flow_rate=[1e-3])]
+    )
+
+    bsn = model.basin.add(
+        Node(2, Point(1, 0)),
+        [basin.Profile(area=1000.0, level=[0.0, 10.0]), basin.State(level=[5.0])],
+    )
+
+    pmp = model.pump.add(
+        Node(3, Point(2, 0)),
+        [
+            pump.Static(
+                control_state=["off", "on"], flow_rate="1e-3", active=[False, True]
+            )
+        ],
+    )
+
+    tmn = model.terminal.add(Node(4, Point(3, 0)))
+
+    dc = model.discrete_control.add(
+        Node(5, Point(1, 1)),
+        [
+            discrete_control.Variable(
+                compound_variable_id=1, listen_node_id=2, variable=["storage"]
+            ),
+            discrete_control.Condition(
+                compound_variable_id=1, condition_id=1, threshold_high=[7500]
+            ),
+            discrete_control.Logic(truth_state=["F", "T"], control_state=["off", "on"]),
+        ],
+    )
+
+    model.link.add(fb, bsn)
+    model.link.add(bsn, pmp)
+    model.link.add(pmp, tmn)
+    model.link.add(dc, pmp)
+
+    return model
+
+
 def connector_node_flow_condition_model() -> Model:
     """DiscreteControl with a condition on the flow through a connector node."""
     model = Model(