Skip to content

Commit

Permalink
Merge branch 'main' into feat/experimental
Browse files Browse the repository at this point in the history
  • Loading branch information
visr authored Oct 21, 2024
2 parents 2f15539 + 4ef4d86 commit b953ee3
Show file tree
Hide file tree
Showing 13 changed files with 190 additions and 66 deletions.
18 changes: 9 additions & 9 deletions core/src/callback.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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)]

Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down
29 changes: 22 additions & 7 deletions core/src/config.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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"
Expand Down
30 changes: 15 additions & 15 deletions core/src/parameter.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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!)}

"""
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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}
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
135 changes: 114 additions & 21 deletions core/src/solve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
9 changes: 9 additions & 0 deletions core/src/util.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
4 changes: 2 additions & 2 deletions core/test/allocation_test.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit b953ee3

Please sign in to comment.