Add step_limiter! to avoid negative flows or too large flows (#1911)

Fixes #1900. Split out of
#1904.

core/src/config.jl

@@ -250,6 +250,18 @@ const algorithms = Dict{String, Type}(
     "Euler" => Euler,
 )
 
+"""
+Check whether the given function has a method that accepts the given kwarg.
+Note that it is possible that methods exist that accept :a and :b individually,
+but not both.
+"""
+function function_accepts_kwarg(f, kwarg)::Bool
+    for method in methods(f)
+        kwarg in Base.kwarg_decl(method) && return true
+    end
+    return false
+end
+
 "Create an OrdinaryDiffEqAlgorithm from solver config"
 function algorithm(solver::Solver; u0 = [])::OrdinaryDiffEqAlgorithm
     algotype = get(algorithms, solver.algorithm, nothing)
@@ -259,19 +271,22 @@ function algorithm(solver::Solver; u0 = [])::OrdinaryDiffEqAlgorithm
             Available options are: ($(options)).")
     end
     kwargs = Dict{Symbol, Any}()
+
     if algotype <: OrdinaryDiffEqNewtonAdaptiveAlgorithm
         kwargs[:nlsolve] = NLNewton(;
             relax = Ribasim.MonitoredBackTracking(; z_tmp = copy(u0), dz_tmp = copy(u0)),
         )
     end
-    # not all algorithms support this keyword
-    kwargs[:autodiff] = solver.autodiff
-    try
-        algotype(; kwargs...)
-    catch
-        pop!(kwargs, :autodiff)
-        algotype(; kwargs...)
+
+    if function_accepts_kwarg(algotype, :step_limiter!)
+        kwargs[:step_limiter!] = Ribasim.limit_flow!
     end
+
+    if function_accepts_kwarg(algotype, :autodiff)
+        kwargs[:autodiff] = solver.autodiff
+    end
+
+    algotype(; kwargs...)
 end
 
 "Convert the saveat Float64 from our Config to SciML's saveat"

core/src/solve.jl

@@ -320,29 +320,14 @@ function formulate_flow!(
 )::Nothing
     (; allocation) = p
     all_nodes_active = p.all_nodes_active[]
-    for (
-        id,
-        inflow_edge,
-        outflow_edge,
-        active,
-        demand_itp,
-        demand,
-        allocated,
-        return_factor,
-        min_level,
-        demand_from_timeseries,
-    ) in zip(
+    for (id, inflow_edge, outflow_edge, active, allocated, return_factor, min_level) in zip(
         user_demand.node_id,
         user_demand.inflow_edge,
         user_demand.outflow_edge,
         user_demand.active,
-        user_demand.demand_itp,
-        # TODO permute these so the nodes are the last dimension, for performance
-        eachrow(user_demand.demand),
         eachrow(user_demand.allocated),
         user_demand.return_factor,
         user_demand.min_level,
-        user_demand.demand_from_timeseries,
     )
         if !(active || all_nodes_active)
             continue
@@ -356,11 +341,7 @@ function formulate_flow!(
         # effectively allocated = demand.
         for priority_idx in eachindex(allocation.priorities)
             alloc_prio = allocated[priority_idx]
-            demand_prio = if demand_from_timeseries
-                demand_itp[priority_idx](t)
-            else
-                demand[priority_idx]
-            end
+            demand_prio = get_demand(user_demand, id, priority_idx, t)
             alloc = min(alloc_prio, demand_prio)
             q += alloc
         end
@@ -718,3 +699,115 @@ function formulate_flows!(
         formulate_flow!(du, user_demand, p, t, current_storage, current_level)
     end
 end
+
+"""
+Clamp the cumulative flow states within the minimum and maximum
+flow rates for the last time step if these flow rate bounds are known.
+"""
+function limit_flow!(
+    u::ComponentVector,
+    integrator::DEIntegrator,
+    p::Parameters,
+    t::Number,
+)::Nothing
+    (; uprev, dt) = integrator
+    (;
+        pump,
+        outlet,
+        linear_resistance,
+        user_demand,
+        tabulated_rating_curve,
+        basin,
+        allocation,
+    ) = p
+
+    # TabulatedRatingCurve flow is in [0, ∞) and can be inactive
+    for (id, active) in zip(tabulated_rating_curve.node_id, tabulated_rating_curve.active)
+        limit_flow!(
+            u.tabulated_rating_curve,
+            uprev.tabulated_rating_curve,
+            id,
+            0.0,
+            Inf,
+            active,
+            dt,
+        )
+    end
+
+    # Pump flow is in [min_flow_rate, max_flow_rate] and can be inactive
+    for (id, min_flow_rate, max_flow_rate, active) in
+        zip(pump.node_id, pump.min_flow_rate, pump.max_flow_rate, pump.active)
+        limit_flow!(u.pump, uprev.pump, id, min_flow_rate, max_flow_rate, active, dt)
+    end
+
+    # Outlet flow is in [min_flow_rate, max_flow_rate] and can be inactive
+    for (id, min_flow_rate, max_flow_rate, active) in
+        zip(outlet.node_id, outlet.min_flow_rate, outlet.max_flow_rate, outlet.active)
+        limit_flow!(u.outlet, uprev.outlet, id, min_flow_rate, max_flow_rate, active, dt)
+    end
+
+    # LinearResistance flow is in [-max_flow_rate, max_flow_rate] and can be inactive
+    for (id, max_flow_rate, active) in zip(
+        linear_resistance.node_id,
+        linear_resistance.max_flow_rate,
+        linear_resistance.active,
+    )
+        limit_flow!(
+            u.linear_resistance,
+            uprev.linear_resistance,
+            id,
+            -max_flow_rate,
+            max_flow_rate,
+            active,
+            dt,
+        )
+    end
+
+    # UserDemand inflow is in [0, total_demand] and can be inactive
+    for (id, active) in zip(user_demand.node_id, user_demand.active)
+        total_demand = sum(
+            get_demand(user_demand, id, priority_idx, t) for
+            priority_idx in eachindex(allocation.priorities)
+        )
+        limit_flow!(
+            u.user_demand_inflow,
+            uprev.user_demand_inflow,
+            id,
+            0.0,
+            total_demand,
+            active,
+            dt,
+        )
+    end
+
+    # Evaporation is in [0, ∞) (a stricter upper bound would require also estimating the area)
+    # Infiltration is in [0, infiltration]
+    for (id, infiltration) in zip(basin.node_id, basin.vertical_flux.infiltration)
+        limit_flow!(u.evaporation, uprev.evaporation, id, 0.0, Inf, true, dt)
+        limit_flow!(u.infiltration, uprev.infiltration, id, 0.0, infiltration, true, dt)
+    end
+
+    return nothing
+end
+
+function limit_flow!(
+    u_component,
+    uprev_component,
+    id::NodeID,
+    min_flow_rate::Number,
+    max_flow_rate::Number,
+    active::Bool,
+    dt::Number,
+)::Nothing
+    u_prev = uprev_component[id.idx]
+    if active
+        u_component[id.idx] = clamp(
+            u_component[id.idx],
+            u_prev + min_flow_rate * dt,
+            u_prev + max_flow_rate * dt,
+        )
+    else
+        u_component[id.idx] = uprev_component[id.idx]
+    end
+    return nothing
+end

core/src/util.jl

@@ -1134,3 +1134,12 @@ function isoutofdomain(u, p, t)
     formulate_storages!(current_storage, u, u, p, t)
     any(<(0), current_storage)
 end
+
+function get_demand(user_demand, id, priority_idx, t)::Float64
+    (; demand_from_timeseries, demand_itp, demand) = user_demand
+    if demand_from_timeseries[id.idx]
+        demand_itp[id.idx][priority_idx](t)
+    else
+        demand[id.idx, priority_idx]
+    end
+end

core/test/run_models_test.jl

@@ -184,6 +184,7 @@ end
     using Logging: Debug, with_logger
     using LoggingExtras
     using SciMLBase: successful_retcode
+    using OrdinaryDiffEqBDF: QNDF
     import Tables
     using Dates
 
@@ -197,11 +198,17 @@ end
     end
 
     @test model isa Ribasim.Model
-    p = model.integrator.p
+
+    (; integrator) = model
+    (; p, alg) = integrator
+
     @test p isa Ribasim.Parameters
     @test isconcretetype(typeof(p))
     @test all(isconcretetype, fieldtypes(typeof(p)))
 
+    @test alg isa QNDF
+    @test alg.step_limiter! == Ribasim.limit_flow!
+
     @test successful_retcode(model)
     @test length(model.integrator.sol) == 2 # start and end
     @test model.integrator.p.basin.current_storage[Float64[]] ≈
@@ -242,7 +249,7 @@ end
     precipitation = model.integrator.p.basin.vertical_flux.precipitation
     @test length(precipitation) == 4
     @test model.integrator.p.basin.current_storage[parent(du)] ≈
-          Float32[720.23611, 694.8785, 415.22371, 1334.3859] atol = 2.0 skip = Sys.isapple()
+          Float32[721.17656, 695.8066, 416.66188, 1334.4879] atol = 2.0 skip = Sys.isapple()
 end
 
 @testitem "Allocation example model" begin