diff --git a/core/src/config.jl b/core/src/config.jl index afad780ed..bf9e38c81 100644 --- a/core/src/config.jl +++ b/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" diff --git a/core/src/solve.jl b/core/src/solve.jl index 9c2555908..3863dc633 100644 --- a/core/src/solve.jl +++ b/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 diff --git a/core/src/util.jl b/core/src/util.jl index f8f17c8c0..77127e99c 100644 --- a/core/src/util.jl +++ b/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 diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl index f033b9f29..bfc634ab1 100644 --- a/core/test/run_models_test.jl +++ b/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