Run allocation at t = 0 (#1496)

Fixes #1389.

core/src/Ribasim.jl

@@ -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,

core/src/allocation_init.jl

@@ -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]))
+    ]
 
     problem[:flow_conservation] = JuMP.@constraint(
         problem,

core/src/model.jl

@@ -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
 
@@ -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)
@@ -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)

core/src/parameter.jl

@@ -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
 """

core/src/util.jl

@@ -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
+    (; u, p, t) = integrator
+    (; allocation, graph, basin) = p
+    (; mean_flows, allocation_models) = allocation
+    (; Δt_allocation) = allocation_models[1]
+    (; 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

core/test/allocation_test.jl

@@ -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(
@@ -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

docs/core/usage.qmd

@@ -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
 --------------| -----

docs/python/examples.ipynb

@@ -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",
@@ -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"
    ]
   },
   {
@@ -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"
    ]
   },
@@ -2114,7 +2117,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,