Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Run allocation at t = 0 #1496

Merged
merged 5 commits into from
May 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion core/src/Ribasim.jl
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ using MetaGraphsNext:
labels,
outneighbor_labels,
inneighbor_labels
using OrdinaryDiffEq: OrdinaryDiffEq, OrdinaryDiffEqRosenbrockAdaptiveAlgorithm
using OrdinaryDiffEq: OrdinaryDiffEq, OrdinaryDiffEqRosenbrockAdaptiveAlgorithm, get_du
using PreallocationTools: DiffCache, get_tmp
using SciMLBase:
init,
Expand Down
6 changes: 5 additions & 1 deletion core/src/allocation_init.jl
Original file line number Diff line number Diff line change
Expand Up @@ -356,7 +356,11 @@ function add_constraints_conservation_node!(
end

# Only the node IDs with conservation constraints on them
node_ids = keys(inflows)
# Discard constraints of the form 0 == 0
node_ids = [
node_id for node_id in keys(inflows) if
!(isempty(inflows[node_id]) && isempty(outflows[node_id]))
]
SouthEndMusic marked this conversation as resolved.
Show resolved Hide resolved

problem[:flow_conservation] = JuMP.@constraint(
problem,
Expand Down
11 changes: 6 additions & 5 deletions core/src/model.jl
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,10 @@ function Model(config::Config)::Model
@show Ribasim.to
end

if config.allocation.use_allocation && is_active(parameters.allocation)
set_initial_allocation_mean_flows!(integrator)
end

return Model(integrator, config, saved)
end

Expand Down Expand Up @@ -207,9 +211,8 @@ function SciMLBase.step!(model::Model, dt::Float64)::Model
# If we are at an allocation time, run allocation before the next physical
# layer timestep. This allows allocation over period (t, t + dt) to use variables
# set over BMI at time t before calling this function.
# Also, don't run allocation at t = 0 since there are no flows yet (#1389).
ntimes = t / config.allocation.timestep
if t > 0 && round(ntimes) ≈ ntimes
if round(ntimes) ≈ ntimes
update_allocation!(integrator)
end
step!(integrator, dt, true)
Expand All @@ -226,10 +229,8 @@ function SciMLBase.solve!(model::Model)::Model
if config.allocation.use_allocation
(; tspan) = integrator.sol.prob
(; timestep) = config.allocation
allocation_times = timestep:timestep:(tspan[end] - timestep)
allocation_times = 0:timestep:(tspan[end] - timestep)
n_allocation_times = length(allocation_times)
# Don't run allocation at t = 0 since there are no flows yet (#1389).
step!(integrator, timestep, true)
for _ in 1:n_allocation_times
update_allocation!(integrator)
step!(integrator, timestep, true)
Expand Down
4 changes: 2 additions & 2 deletions core/src/parameter.jl
Original file line number Diff line number Diff line change
Expand Up @@ -509,14 +509,14 @@ outflow_edge: outgoing flow edge metadata
The ID of the source node is always the ID of the UserDemand node
active: whether this node is active and thus demands water
realized_bmi: Cumulative inflow volume, for read or reset by BMI only
demand: water flux demand of UserDemand per priority over time
demand: water flux demand of UserDemand per priority (node_idx, priority_idx)
Each UserDemand has a demand for all priorities,
which is 0.0 if it is not provided explicitly.
demand_reduced: the total demand reduced by allocated flows. This is used for goal programming,
and requires separate memory from `demand` since demands can come from the BMI
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
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
"""
Expand Down
30 changes: 30 additions & 0 deletions core/src/util.jl
Original file line number Diff line number Diff line change
Expand Up @@ -750,3 +750,33 @@ function get_basin_idx(edge_metadata::EdgeMetadata, id::NodeID)::Int32
0
end
end

"""
We want to perform allocation at t = 0 but there are no mean flows available
as input. Therefore we set the instantaneous flows as the mean flows as allocation input.
"""
function set_initial_allocation_mean_flows!(integrator)::Nothing
SouthEndMusic marked this conversation as resolved.
Show resolved Hide resolved
(; u, p, t) = integrator
(; allocation, graph, basin) = p
(; mean_flows, allocation_models) = allocation
(; Δt_allocation) = allocation_models[1]
SouthEndMusic marked this conversation as resolved.
Show resolved Hide resolved
(; vertical_flux) = basin
vertical_flux = get_tmp(vertical_flux, 0)

# At the time of writing water_balance! already
# gets called once at the problem initialization, this
# one is just to make sure.
water_balance!(get_du(integrator), u, p, t)

for (edge, value) in mean_flows
if edge[1] == edge[2]
q = get_influx(basin, edge[1])
else
q = get_flow(graph, edge..., 0)
end
# Multiply by Δt_allocation as averaging divides by this factor
# in update_allocation!
value[] = q * Δt_allocation
end
return nothing
end
58 changes: 29 additions & 29 deletions core/test/allocation_test.jl
Original file line number Diff line number Diff line change
Expand Up @@ -323,17 +323,23 @@ end

@testitem "Allocation level control" begin
import JuMP
using Ribasim: NodeID

toml_path = normpath(@__DIR__, "../../generated_testmodels/level_demand/ribasim.toml")
@test ispath(toml_path)
model = Ribasim.run(toml_path)

storage = Ribasim.get_storages_and_levels(model).storage[1, :]
t = Ribasim.tsaves(model)
model = Ribasim.Model(toml_path)

p = model.integrator.p
(; user_demand, graph, allocation, basin, level_demand) = p

# Initial "integrated" vertical flux
@test allocation.mean_flows[(NodeID(:Basin, 2), NodeID(:Basin, 2))][] ≈ 1e2

Ribasim.solve!(model)

storage = Ribasim.get_storages_and_levels(model).storage[1, :]
t = Ribasim.tsaves(model)

d = user_demand.demand_itp[1][2](0)
ϕ = 1e-3 # precipitation
q = Ribasim.get_flow(
Expand All @@ -346,48 +352,42 @@ end
l_max = level_demand.max_level[1](0)
Δt_allocation = allocation.allocation_models[1].Δt_allocation

# Until the first allocation solve, the UserDemand abstracts fully
stage_1 = t .<= Δt_allocation
u_stage_1(τ) = storage[1] + (q + ϕ - d) * τ
@test storage[stage_1] ≈ u_stage_1.(t[stage_1]) rtol = 1e-4

# In this section the Basin leaves no supply for the UserDemand
stage_2 = Δt_allocation .<= t .<= 3 * Δt_allocation
stage_2_start_idx = findfirst(stage_2)
u_stage_2(τ) = storage[stage_2_start_idx] + (q + ϕ) * (τ - t[stage_2_start_idx])
@test storage[stage_2] ≈ u_stage_2.(t[stage_2]) rtol = 1e-4
stage_1 = t .<= 2 * Δt_allocation
u_stage_1(τ) = storage[1] + (q + ϕ) * τ
@test storage[stage_1] ≈ u_stage_1.(t[stage_1]) rtol = 1e-4

# In this section (and following sections) the basin has no longer a (positive) demand,
# since precipitation provides enough water to get the basin to its target level
# The FlowBoundary flow gets fully allocated to the UserDemand
stage_3 = 3 * Δt_allocation .<= t .<= 8 * Δt_allocation
stage_3_start_idx = findfirst(stage_3)
u_stage_3(τ) = storage[stage_3_start_idx] + ϕ * (τ - t[stage_3_start_idx])
@test storage[stage_3] ≈ u_stage_3.(t[stage_3]) rtol = 1e-4
stage_2 = 2 * Δt_allocation .<= t .<= 8 * Δt_allocation
stage_2_start_idx = findfirst(stage_2)
u_stage_2(τ) = storage[stage_2_start_idx] + ϕ * (τ - t[stage_2_start_idx])
@test storage[stage_2] ≈ u_stage_2.(t[stage_2]) rtol = 1e-4

# In this section the basin enters its surplus stage,
# even though initially the level is below the maximum level. This is because the simulation
# anticipates that the current precipitation is going to bring the basin level over
# its maximum level
stage_4 = 8 * Δt_allocation .<= t .<= 12 * Δt_allocation
stage_4_start_idx = findfirst(stage_4)
u_stage_4(τ) = storage[stage_4_start_idx] + (q + ϕ - d) * (τ - t[stage_4_start_idx])
@test storage[stage_4] ≈ u_stage_4.(t[stage_4]) rtol = 1e-4
stage_3 = 8 * Δt_allocation .<= t .<= 13 * Δt_allocation
stage_3_start_idx = findfirst(stage_3)
u_stage_3(τ) = storage[stage_3_start_idx] + (q + ϕ - d) * (τ - t[stage_3_start_idx])
@test storage[stage_3] ≈ u_stage_3.(t[stage_3]) rtol = 1e-4

# At the start of this section precipitation stops, and so the UserDemand
# partly uses surplus water from the basin to fulfill its demand
stage_5 = 13 * Δt_allocation .<= t .<= 16 * Δt_allocation
stage_5_start_idx = findfirst(stage_5)
u_stage_5(τ) = storage[stage_5_start_idx] + (q - d) * (τ - t[stage_5_start_idx])
@test storage[stage_5] ≈ u_stage_5.(t[stage_5]) rtol = 1e-4
stage_4 = 13 * Δt_allocation .<= t .<= 17 * Δt_allocation
stage_4_start_idx = findfirst(stage_4)
u_stage_4(τ) = storage[stage_4_start_idx] + (q - d) * (τ - t[stage_4_start_idx])
@test storage[stage_4] ≈ u_stage_4.(t[stage_4]) rtol = 1e-4

# From this point the basin is in a dynamical equilibrium,
# since the basin has no supply so the UserDemand abstracts precisely
# the flow from the level boundary
stage_6 = 17 * Δt_allocation .<= t
stage_6_start_idx = findfirst(stage_6)
u_stage_6(τ) = storage[stage_6_start_idx]
@test storage[stage_6] ≈ u_stage_6.(t[stage_6]) rtol = 1e-4
stage_5 = 18 * Δt_allocation .<= t
stage_5_start_idx = findfirst(stage_5)
u_stage_5(τ) = storage[stage_5_start_idx]
@test storage[stage_5] ≈ u_stage_5.(t[stage_5]) rtol = 1e-4

# Isolated LevelDemand + Basin pair to test optional min_level
problem = allocation.allocation_models[2].problem
Expand Down
1 change: 1 addition & 0 deletions docs/core/usage.qmd
Original file line number Diff line number Diff line change
Expand Up @@ -847,6 +847,7 @@ control_state | String
## Allocation - `allocation.arrow`

The allocation table contains a record of allocation results: when it happened, for which node, in which allocation network, and what the demand, allocated flow and realized flow were.
The realized values at the starting time of the simulation can be ignored.

column | type
--------------| -----
Expand Down
21 changes: 12 additions & 9 deletions docs/python/examples.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -1356,7 +1356,9 @@
"\n",
"df_allocation_wide[\"demand\"].plot(ax=axs[0], ls=\":\")\n",
"df_allocation_wide[\"allocated\"].plot(ax=axs[1], ls=\"--\")\n",
"df_allocation_wide[\"realized\"].plot(ax=axs[2])\n",
"df_allocation_wide.xs(1, level=\"priority\", axis=1)[\"realized\"].plot(\n",
" ax=axs[2], color=[\"C0\", \"C2\", \"C3\"]\n",
")\n",
"\n",
"fig.tight_layout()\n",
"loc = plticker.MultipleLocator(2)\n",
Expand All @@ -1375,9 +1377,11 @@
"source": [
"Some things to note about this plot:\n",
"\n",
"- Abstraction behaves somewhat erratically at the start of the simulation. This is because allocation is based on flows computed in the physical layer, and at the start of the simulation these are not known yet.\n",
"- The realized flow at the start time is not correct, as there is no realized flow yet. The given value is how much of its total demand a user can abstract in the physical layer.\n",
"\n",
"- No flow was allocated to UserDemand 13 so that is not plotted.\n",
"\n",
"- Although there is a plotted line for abstraction per priority, abstraction is actually accumulated over all priorities per user.\n"
"- Abstraction is accumulated over all priorities per user.\n"
]
},
{
Expand Down Expand Up @@ -1655,11 +1659,10 @@
"source": [
"Stages:\n",
"\n",
"- In the first stage the UserDemand abstracts fully, so the net change of Basin #2 is $q + \\phi - d$;\n",
"- In the second stage the Basin takes precedence so the UserDemand doesn't abstract, hence the net change of Basin #2 is $q + \\phi$;\n",
"- In the third stage (and following stages) the Basin no longer has a positive demand, since precipitation provides enough water to get the Basin to its target level. The FlowBoundary flow gets fully allocated to the UserDemand, hence the net change of Basin #2 is $\\phi$;\n",
"- In the fourth stage the Basin enters its surplus stage, even though initially the level is below the maximum level. This is because the simulation anticipates that the current precipitation is going to bring the Basin level over its maximum level. The net change of Basin #2 is now $q + \\phi - d$;\n",
"- At the start of the fifth stage the precipitation stops, and so the UserDemand partly uses surplus water from the Basin to fulfill its demand. The net change of Basin #2 becomes $q - d$.\n",
"- In the first stage the Basin takes precedence so the UserDemand doesn't abstract, hence the net change of Basin #2 is $q + \\phi$;\n",
"- In the second stage (and following stages) the Basin no longer has a positive demand, since precipitation provides enough water to get the Basin to its target level. The FlowBoundary flow gets fully allocated to the UserDemand, hence the net change of Basin #2 is $\\phi$;\n",
"- In the third stage the Basin enters its surplus stage, even though initially the level is below the maximum level. This is because the simulation anticipates that the current precipitation is going to bring the Basin level over its maximum level. The net change of Basin #2 is now $q + \\phi - d$;\n",
"- At the start of the fourth stage the precipitation stops, and so the UserDemand partly uses surplus water from the Basin to fulfill its demand. The net change of Basin #2 becomes $q - d$.\n",
"- In the final stage the Basin is in a dynamical equilibrium, since the Basin has no supply so the user abstracts precisely the flow from the LevelBoundary.\n"
]
},
Expand Down Expand Up @@ -2114,7 +2117,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.2"
"version": "3.12.3"
}
},
"nbformat": 4,
Expand Down