Merge branch 'main' into feat/experimental

core/src/callback.jl

@@ -109,8 +109,8 @@ Update with the latest timestep:
 
 During these cumulative flow updates, we can also update the mass balance of the system,
 as each flow carries mass, based on the concentrations of the flow source.
-Specifically, we first use all the inflows to update the mass of the basins, recalculate
-the basin concentration(s) and then remove the mass that is being lost to the outflows.
+Specifically, we first use all the inflows to update the mass of the Basins, recalculate
+the Basin concentration(s) and then remove the mass that is being lost to the outflows.
 """
 function update_cumulative_flows!(u, t, integrator)::Nothing
     (; p, uprev, tprev, dt) = integrator
@@ -130,7 +130,7 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
 
     # Reset cumulative flows, used to calculate the concentration
     # of the basins after processing inflows only
-    fill!(basin.cumulative_in, 0.0)
+    basin.cumulative_in .= 0.0
 
     # Update cumulative forcings which are integrated exactly
     @. basin.cumulative_drainage += vertical_flux.drainage * dt
@@ -264,10 +264,10 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
         end
     end
 
-    # Update the basin concentrations based on the added mass and flows
+    # Update the Basin concentrations based on the added mass and flows
     basin.concentration_state .= basin.mass ./ (basin.storage_prev .+ basin.cumulative_in)
 
-    # Process all mass outflows from basins
+    # Process all mass outflows from Basins
     for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge)
         from_node = inflow_edge.edge[1]
         to_node = outflow_edge.edge[2]
@@ -311,7 +311,7 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
         error("Negative mass(es) detected")
     end
 
-    # Update the basin concentrations again based on the removed mass
+    # Update the Basin concentrations again based on the removed mass
     basin.concentration_state .= basin.mass ./ basin.current_storage[parent(u)]
     basin.storage_prev .= basin.current_storage[parent(u)]
 
@@ -320,8 +320,8 @@ end
 
 """
 Given an edge (from_id, to_id), compute the cumulative flow over that
-edge over the latest timestep. If from_id and to_id are both the same basin,
-the function returns the sum of the basin forcings.
+edge over the latest timestep. If from_id and to_id are both the same Basin,
+the function returns the sum of the Basin forcings.
 """
 function flow_update_on_edge(
     integrator::DEIntegrator,
@@ -366,7 +366,7 @@ end
 """
 Save all cumulative forcings and flows over edges over the latest timestep,
 Both computed by the solver and integrated exactly. Also computes the total horizontal
-inflow and outflow per basin.
+inflow and outflow per Basin.
 """
 function save_flow(u, t, integrator)
     (; p) = integrator

core/src/config.jl

@@ -264,6 +264,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)
@@ -273,19 +285,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/parameter.jl

@@ -106,7 +106,7 @@ const ScalarInterpolation = LinearInterpolation{
     Float64,
 }
 
-set_zero!(v) = fill!(v, zero(eltype(v)))
+set_zero!(v) = v .= zero(eltype(v))
 const Cache = LazyBufferCache{Returns{Int}, typeof(set_zero!)}
 
 """
@@ -266,12 +266,12 @@ abstract type AbstractDemandNode <: AbstractParameterNode end
 In-memory storage of saved mean flows for writing to results.
 
 - `flow`: The mean flows on all edges and state-dependent forcings
-- `inflow`: The sum of the mean flows coming into each basin
-- `outflow`: The sum of the mean flows going out of each basin
+- `inflow`: The sum of the mean flows coming into each Basin
+- `outflow`: The sum of the mean flows going out of each Basin
 - `flow_boundary`: The exact integrated mean flows of flow boundaries
 - `precipitation`: The exact integrated mean precipitation
 - `drainage`: The exact integrated mean drainage
-- `concentration`: Concentrations for each basin and substance
+- `concentration`: Concentrations for each Basin and substance
 - `balance_error`: The (absolute) water balance error
 - `relative_error`: The relative water balance error
 - `t`: Endtime of the interval over which is averaged
@@ -356,15 +356,15 @@ end
     # Concentrations
     # Config setting to enable/disable evaporation of mass
     evaporate_mass::Bool = true
-    # Cumulative inflow for each basin at a given time
+    # Cumulative inflow for each Basin at a given time
     cumulative_in::Vector{Float64} = zeros(length(node_id))
-    # Storage for each basin at the previous time step
+    # Storage for each Basin at the previous time step
     storage_prev::Vector{Float64} = zeros(length(node_id))
-    # matrix with concentrations for each basin and substance
-    concentration_state::Matrix{Float64}  # basin, substance
-    # matrix with boundary concentrations for each boundary, basin and substance
+    # matrix with concentrations for each Basin and substance
+    concentration_state::Matrix{Float64}  # Basin, substance
+    # matrix with boundary concentrations for each boundary, Basin and substance
     concentration::Array{Float64, 3}
-    # matrix with mass for each basin and substance
+    # matrix with mass for each Basin and substance
     mass::Matrix{Float64}
     # substances in use by the model (ordered like their axis in the concentration matrices)
     substances::OrderedSet{Symbol}
@@ -482,8 +482,8 @@ end
 """
 node_id: node ID of the LevelBoundary node
 active: whether this node is active
-level: the fixed level of this 'infinitely big basin'
-concentration: matrix with boundary concentrations for each basin and substance
+level: the fixed level of this 'infinitely big Basin'
+concentration: matrix with boundary concentrations for each Basin and substance
 concentration_time: Data source for concentration updates
 """
 @kwdef struct LevelBoundary{C} <: AbstractParameterNode
@@ -501,7 +501,7 @@ active: whether this node is active and thus contributes flow
 cumulative_flow: The exactly integrated cumulative boundary flow since the start of the simulation
 cumulative_flow_saveat: The exactly integrated cumulative boundary flow since the last saveat
 flow_rate: flow rate (exact)
-concentration: matrix with boundary concentrations for each basin and substance
+concentration: matrix with boundary concentrations for each Basin and substance
 concentration_time: Data source for concentration updates
 """
 @kwdef struct FlowBoundary{C} <: AbstractParameterNode
@@ -776,8 +776,8 @@ demand_itp: Timeseries interpolation objects for demands
 demand_from_timeseries: If false the demand comes from the BMI or is fixed
 allocated: water flux currently allocated to UserDemand per priority (node_idx, priority_idx)
 return_factor: the factor in [0,1] of how much of the abstracted water is given back to the system
-min_level: The level of the source basin below which the UserDemand does not abstract
-concentration: matrix with boundary concentrations for each basin and substance
+min_level: The level of the source Basin below which the UserDemand does not abstract
+concentration: matrix with boundary concentrations for each Basin and substance
 concentration_time: Data source for concentration updates
 """
 @kwdef struct UserDemand{C} <: AbstractDemandNode

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/allocation_test.jl

@@ -582,8 +582,8 @@ end
     (; allocation_models) = p.allocation
     (; basin, level_demand, graph) = p
 
-    fill!(level_demand.max_level[1].u, Inf)
-    fill!(level_demand.max_level[2].u, Inf)
+    level_demand.max_level[1].u .= Inf
+    level_demand.max_level[2].u .= Inf
 
     # Given a max_level of Inf, the basin capacity is 0.0 because it is not possible for the basin level to be > Inf
     @test Ribasim.get_basin_capacity(allocation_models[1], u, p, t, basin.node_id[1]) == 0.0