From 9238b0c1351fb5fd2c61a7aedab5295b630494b7 Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Tue, 12 Sep 2023 15:15:00 +0200
Subject: [PATCH] clamp reduction_factor and add it to TabulatedRatingCurve
 (#587)

If the rating curve prescribes large Q even when the upstream Basin is
nearly empty, this leads to an incorrect water balance.

Replaces #562.
---
 core/src/solve.jl       | 21 ++++++++++++++-------
 core/test/run_models.jl |  2 +-
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/core/src/solve.jl b/core/src/solve.jl
index 227ced273..6e28a7f32 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -520,7 +520,7 @@ function formulate!(
         bottom = basin.level[i][1]
         fixed_area = basin.area[i][end]
         depth = max(level - bottom, 0.0)
-        reduction_factor = min(depth, 0.1) / 0.1
+        reduction_factor = clamp(depth, 0.0, 0.1) / 0.1
 
         precipitation = fixed_area * basin.precipitation[i]
         evaporation = area * reduction_factor * basin.potential_evaporation[i]
@@ -594,7 +594,7 @@ function continuous_control!(
             controlled_node_idx = findsorted(pump.node_id, controlled_node_id)
 
             listened_basin_storage = u.storage[listened_node_idx]
-            reduction_factor = min(listened_basin_storage, 10.0) / 10.0
+            reduction_factor = clamp(listened_basin_storage, 0.0, 10.0) / 10.0
         else
             controlled_node_idx = findsorted(outlet.node_id, controlled_node_id)
 
@@ -603,7 +603,7 @@ function continuous_control!(
             has_index, upstream_basin_idx = id_index(basin.node_id, upstream_node_id)
             if has_index
                 upstream_basin_storage = u.storage[upstream_basin_idx]
-                reduction_factor = min(upstream_basin_storage, 10.0) / 10.0
+                reduction_factor = clamp(upstream_basin_storage, 0.0, 10.0) / 10.0
             else
                 reduction_factor = 1.0
             end
@@ -736,11 +736,12 @@ Directed graph: outflow is positive!
 function formulate_flow!(
     tabulated_rating_curve::TabulatedRatingCurve,
     p::Parameters,
+    storage::AbstractVector,
     current_level::AbstractVector,
     flow::AbstractMatrix,
     t::Float64,
 )::Nothing
-    (; connectivity) = p
+    (; basin, connectivity) = p
     (; graph_flow) = connectivity
     (; node_id, active, tables) = tabulated_rating_curve
     for (i, id) in enumerate(node_id)
@@ -748,7 +749,13 @@ function formulate_flow!(
         downstream_ids = outneighbors(graph_flow, id)
 
         if active[i]
-            q = tables[i](get_level(p, upstream_basin_id, current_level, t))
+            hasindex, basin_idx = id_index(basin.node_id, upstream_basin_id)
+            @assert hasindex "TabulatedRatingCurve must be downstream of a Basin"
+            s = storage[basin_idx]
+            reduction_factor = clamp(s, 0.0, 10.0) / 10.0
+            q =
+                reduction_factor *
+                tables[i](get_level(p, upstream_basin_id, current_level, t))
         else
             q = 0.0
         end
@@ -931,7 +938,7 @@ function formulate_flow!(
         if hasindex
             # Pumping from basin
             s = storage[basin_idx]
-            reduction_factor = min(s, 10.0) / 10.0
+            reduction_factor = clamp(s, 0.0, 10.0) / 10.0
             q = reduction_factor * rate
         else
             # Pumping from level boundary
@@ -984,7 +991,7 @@ function formulate_flows!(
 
     formulate_flow!(linear_resistance, p, current_level, flow, t)
     formulate_flow!(manning_resistance, p, current_level, flow, t)
-    formulate_flow!(tabulated_rating_curve, p, current_level, flow, t)
+    formulate_flow!(tabulated_rating_curve, p, storage, current_level, flow, t)
     formulate_flow!(flow_boundary, p, flow, t)
     formulate_flow!(fractional_flow, flow, p)
     formulate_flow!(pump, p, flow, storage)
diff --git a/core/test/run_models.jl b/core/test/run_models.jl
index 466bbe2bc..1658824de 100644
--- a/core/test/run_models.jl
+++ b/core/test/run_models.jl
@@ -86,7 +86,7 @@ end
     model = Ribasim.run(toml_path)
     @test model isa Ribasim.Model
     @test successful_retcode(model)
-    @test model.integrator.sol.u[end] ≈ Float32[5.951445, 727.9898] skip = Sys.isapple()
+    @test model.integrator.sol.u[end] ≈ Float32[8.41143, 725.5068] skip = Sys.isapple()
     # the highest level in the dynamic table is updated to 1.2 from the callback
     @test model.integrator.p.tabulated_rating_curve.tables[end].t[end] == 1.2
 end