Allocate to UserDemand from directly connected basin if possible (#1581)

Fixes #1545

Some notes on this PR:

- This still only works when the `Basin` upstream of a `UserDemand` has
a level demand to define the source capacity of the basin
- It makes sense that this flow from the `Basin` to the `UserDemand`
also shows up in `allocation_flow.arrow`, which requires some
refactoring. That is fine though, as that refactor is needed anyway for
#565, where we probably want
to aggregate the allocation flows for output over the optimizations for
all sources in stead of saving them per source
- I'm not sure how to directly test this

---------

Co-authored-by: Martijn Visser <[email protected]>

core/src/allocation_init.jl

@@ -492,6 +492,7 @@ function AllocationModel(
 )::AllocationModel
     capacity = get_capacity(p, subnetwork_id)
     problem = allocation_problem(p, capacity, subnetwork_id)
+    flow_priority = JuMP.Containers.SparseAxisArray(Dict(only(problem[:F].axes) .=> 0.0))
 
-    return AllocationModel(; subnetwork_id, capacity, problem, Δt_allocation)
+    return AllocationModel(; subnetwork_id, capacity, flow_priority, problem, Δt_allocation)
 end

core/src/allocation_optim.jl

@@ -461,7 +461,7 @@ function set_initial_demands_level!(
 )::Nothing
     (; subnetwork_id, problem) = allocation_model
     (; graph, basin) = p
-    (; node_id, demand) = basin
+    (; demand) = basin
 
     node_ids_level_demand = only(problem[:basin_outflow].axes)
 
@@ -558,7 +558,7 @@ function adjust_demands!(
     ::LevelDemand,
 )::Nothing
     (; graph, basin) = p
-    (; node_id, demand) = basin
+    (; demand) = basin
     (; subnetwork_id, problem) = allocation_model
     F_basin_in = problem[:F_basin_in]
 
@@ -780,10 +780,9 @@ function save_allocation_flows!(
     priority::Int32,
     optimization_type::OptimizationType.T,
 )::Nothing
-    (; problem, subnetwork_id, capacity) = allocation_model
+    (; flow_priority, problem, subnetwork_id, capacity) = allocation_model
     (; allocation, graph) = p
     (; record_flow) = allocation
-    F = problem[:F]
     F_basin_in = problem[:F_basin_in]
     F_basin_out = problem[:F_basin_out]
 
@@ -802,12 +801,12 @@ function save_allocation_flows!(
         flow_rate = 0.0
 
         if haskey(graph, edge_1...)
-            flow_rate += JuMP.value(F[edge_1])
+            flow_rate += flow_priority[edge_1]
             sign_2 = -1.0
             edge_metadata = graph[edge_1...]
         else
             edge_1_reverse = reverse(edge_1)
-            flow_rate -= JuMP.value(F[edge_1_reverse])
+            flow_rate -= flow_priority[edge_1_reverse]
             sign_2 = 1.0
             edge_metadata = graph[edge_1_reverse...]
         end
@@ -817,7 +816,7 @@ function save_allocation_flows!(
         if edge_2 == reverse(edge_1) &&
            !(edge_1[1].type == NodeType.UserDemand || edge_1[2].type == NodeType.UserDemand)
             # If so, these edges are both processed in this iteration
-            flow_rate += sign_2 * JuMP.value(F[edge_2])
+            flow_rate += sign_2 * flow_priority[edge_2]
             skip = true
         end
 
@@ -857,6 +856,45 @@ function save_allocation_flows!(
     return nothing
 end
 
+function allocate_to_users_from_connected_basin!(
+    allocation_model::AllocationModel,
+    p::Parameters,
+    priority_idx::Int,
+)::Nothing
+    (; flow_priority, problem) = allocation_model
+    (; graph, user_demand) = p
+
+    # Get all UserDemand nodes from this subnetwork
+    node_ids_user_demand = only(problem[:source_user].axes)
+    for node_id in node_ids_user_demand
+
+        # Check whether the upstream basin has a level demand
+        # and thus can act as a source
+        upstream_basin_id = user_demand.inflow_edge[node_id.idx].edge[1]
+        if has_external_demand(graph, upstream_basin_id, :level_demand)[1]
+
+            # The demand of the UserDemand node at the current priority
+            demand = user_demand.demand_reduced[node_id.idx, priority_idx]
+
+            # The capacity of the upstream basin
+            constraint = problem[:basin_outflow][upstream_basin_id]
+            capacity = JuMP.normalized_rhs(constraint)
+
+            # The allocated amount
+            allocated = min(demand, capacity)
+
+            # Subtract the allocated amount from the user demand and basin capacity
+            user_demand.demand_reduced[node_id.idx, priority_idx] -= allocated
+            JuMP.set_normalized_rhs(constraint, capacity - allocated)
+
+            # Add the allocated flow
+            flow_priority[(upstream_basin_id, node_id)] += allocated
+        end
+    end
+
+    return nothing
+end
+
 function optimize_priority!(
     allocation_model::AllocationModel,
     u::ComponentVector,
@@ -865,10 +903,15 @@ function optimize_priority!(
     priority_idx::Int,
     optimization_type::OptimizationType.T,
 )::Nothing
-    (; problem) = allocation_model
+    (; problem, flow_priority) = allocation_model
     (; allocation) = p
     (; priorities) = allocation
 
+    # Start the values of the flows at this priority at 0.0
+    for edge in keys(flow_priority.data)
+        flow_priority[edge] = 0.0
+    end
+
     set_capacities_flow_demand_outflow!(allocation_model, p, priority_idx)
 
     # Set the objective depending on the demands
@@ -878,6 +921,10 @@ function optimize_priority!(
     # https://jump.dev/JuMP.jl/v1.16/manual/objective/#Modify-an-objective-coefficient
     set_objective_priority!(allocation_model, p, u, t, priority_idx)
 
+    # Allocate to UserDemand nodes from the directly connected basin
+    # This happens outside the JuMP optimization
+    allocate_to_users_from_connected_basin!(allocation_model, p, priority_idx)
+
     # Solve the allocation problem for this priority
     JuMP.optimize!(problem)
     @debug JuMP.solution_summary(problem)
@@ -889,6 +936,11 @@ function optimize_priority!(
         )
     end
 
+    # Add the values of the flows at this priority
+    for edge in only(problem[:F].axes)
+        flow_priority[edge] += JuMP.value(problem[:F][edge])
+    end
+
     # Assign the allocations to the UserDemand for this priority
     assign_allocations!(allocation_model, p, priority_idx, optimization_type)
 

core/src/callback.jl

@@ -451,7 +451,7 @@ end
 "Solve the allocation problem for all demands and assign allocated abstractions."
 function update_allocation!(integrator)::Nothing
     (; p, t, u) = integrator
-    (; allocation, basin) = p
+    (; allocation) = p
     (; allocation_models, mean_input_flows, mean_realized_flows) = allocation
 
     # Don't run the allocation algorithm if allocation is not active

core/src/parameter.jl

@@ -80,12 +80,14 @@ Store information for a subnetwork used for allocation.
 
 subnetwork_id: The ID of this allocation network
 capacity: The capacity per edge of the allocation network, as constrained by nodes that have a max_flow_rate
+flow_priority: The flows over all the edges in the subnetwork for a certain priority (used for allocation_flow output)
 problem: The JuMP.jl model for solving the allocation problem
 Δt_allocation: The time interval between consecutive allocation solves
 """
 @kwdef struct AllocationModel
     subnetwork_id::Int32
     capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}}
+    flow_priority::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}}
     problem::JuMP.Model
     Δt_allocation::Float64
 end

core/test/allocation_test.jl

@@ -283,7 +283,7 @@ end
     # 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_2 = 2 * Δt_allocation .<= t .<= 8 * Δt_allocation
+    stage_2 = 2 * Δt_allocation .<= t .<= 9 * Δ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
@@ -292,22 +292,22 @@ end
     # 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_3 = 8 * Δt_allocation .<= t .<= 13 * Δt_allocation
+    stage_3 = 9 * Δ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_4 = 13 * Δt_allocation .<= t .<= 17 * Δt_allocation
+    stage_4 = 13 * Δt_allocation .<= t .<= 15 * Δ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_5 = 18 * Δt_allocation .<= t
+    stage_5 = 16 * Δ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
@@ -551,6 +551,28 @@ end
     @test all(isapprox.(fractions[1], fractions[4], atol = 1e-4))
 end
 
+@testitem "direct_basin_allocation" begin
+    using Ribasim: NodeID
+    import SQLite
+    import JuMP
+
+    toml_path = normpath(@__DIR__, "../../generated_testmodels/level_demand/ribasim.toml")
+    model = Ribasim.Model(toml_path)
+    (; p) = model.integrator
+    t = 0.0
+    u = model.integrator.u
+    priority_idx = 2
+
+    allocation_model = first(p.allocation.allocation_models)
+    Ribasim.set_initial_values!(allocation_model, p, u, t)
+    Ribasim.set_objective_priority!(allocation_model, p, u, t, priority_idx)
+    Ribasim.allocate_to_users_from_connected_basin!(allocation_model, p, priority_idx)
+    flow_data = allocation_model.flow_priority.data
+    @test flow_data[(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] == 0.0
+    @test flow_data[(NodeID(:Basin, 2, p), NodeID(:UserDemand, 3, p))] == 0.0015
+    @test flow_data[(NodeID(:UserDemand, 3, p), NodeID(:Basin, 5, p))] == 0.0
+end
+
 @testitem "level_demand_without_max_level" begin
     using Ribasim: NodeID, get_basin_capacity, outflow_id
     using JuMP

docs/concept/allocation.qmd

@@ -7,6 +7,8 @@ Allocation is the process of assigning an allocated flow rate to demand nodes in
 
 The allocation problem is solved per subnetwork (and main network) of the Ribasim model. Each subnetwork is used to formulate an optimization problem with the [JuMP](https://jump.dev/JuMP.jl/stable/) package, which is solved using the [HiGHS solver](https://highs.dev/). For more in-depth information see also the example of solving the maximum flow problem with `JuMP.jl` [here](https://jump.dev/JuMP.jl/stable/tutorials/linear/network_flows/#The-max-flow-problem).
 
+Before the optimization for each priority there is a simple step that tries to allocate flow to the `UserDemand` nodes from the their directly connected basin.
+
 :::{.callout-note}
 within this *Allocation* section the main network is also considered to be a subnetwork.
 :::

docs/reference/node/user-demand.qmd

@@ -4,8 +4,8 @@ title: "UserDemand"
 
 A UserDemand takes water from the Basin that supplies it.
 
-When allocation is not used, UserDemand attempts to extract the full demand from the Basin.
-When allocation is used, water is allocated based on its priority, where priority 1 denotes the most important demand.
+When allocation is not used, a `UserDemand` node attempts to extract the full demand from the connected Basin.
+When allocation is used, the amount a `UserDemand` node is allowed to abstract is determined by the [allocation algorithm](../../concept/allocation.qmd). This algorithm first tries to allocate from the directly connected basin, and then from other sources whose flow can reach the `UserDemand` node.
 
 When the connected Basin is almost empty or reaches the minimum level at which the UserDemand can extract water (`min_level`), it will stop extraction.