From 03c60ccdc782a453c3e06c6cf6740f74124b14c7 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 25 Jun 2024 14:02:14 +0200
Subject: [PATCH 1/8] Allocate to UserDemand from directly connected basin if
 possible

---
 core/src/allocation_optim.jl | 44 ++++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index bb520208e..1723d514f 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -457,7 +457,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)
 
@@ -554,7 +554,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]
 
@@ -853,6 +853,42 @@ function save_allocation_flows!(
     return nothing
 end
 
+function allocate_to_users_from_connected_basin!(
+    allocation_model::AllocationModel,
+    p::Parameters,
+    priority_idx::Int,
+)::Nothing
+    (; problem) = allocation_model
+    (; graph, user_demand) = p
+
+    # Get al 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)
+        end
+    end
+
+    return nothing
+end
+
 function optimize_priority!(
     allocation_model::AllocationModel,
     u::ComponentVector,
@@ -874,6 +910,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)

From d755c58be34bff40ca740a806778656e7a81759b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 25 Jun 2024 14:34:16 +0200
Subject: [PATCH 2/8] Fix test

---
 core/test/allocation_test.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index a62791496..bc6328117 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -365,7 +365,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
@@ -374,14 +374,14 @@ 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
@@ -389,7 +389,7 @@ end
     # 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

From 2fb6a2a92dd342ea07dd82c4a082295fe5875ace Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 25 Jun 2024 16:36:22 +0200
Subject: [PATCH 3/8] Add directly allocated flow to flow output

---
 core/src/allocation_init.jl  |  3 ++-
 core/src/allocation_optim.jl | 32 +++++++++++++++++++++++++-------
 core/src/callback.jl         |  2 +-
 core/src/parameter.jl        |  2 ++
 4 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl
index 20cd9eb06..8e0dec8dd 100644
--- a/core/src/allocation_init.jl
+++ b/core/src/allocation_init.jl
@@ -543,6 +543,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
diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 1723d514f..3bd348c78 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -776,10 +776,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]
 
@@ -798,12 +797,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
@@ -813,7 +812,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
 
@@ -855,6 +854,7 @@ end
 
 function allocate_to_users_from_connected_basin!(
     allocation_model::AllocationModel,
+    flow_priority,
     p::Parameters,
     priority_idx::Int,
 )::Nothing
@@ -883,6 +883,9 @@ function allocate_to_users_from_connected_basin!(
             # 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
 
@@ -897,10 +900,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
@@ -912,7 +920,12 @@ function optimize_priority!(
 
     # 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)
+    allocate_to_users_from_connected_basin!(
+        allocation_model,
+        flow_priority,
+        p,
+        priority_idx,
+    )
 
     # Solve the allocation problem for this priority
     JuMP.optimize!(problem)
@@ -925,6 +938,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)
 
diff --git a/core/src/callback.jl b/core/src/callback.jl
index be2f14f06..56b64afae 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -488,7 +488,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
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 3a1aa523c..f0eb1cbf9 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -79,12 +79,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

From 34ca8a113e97de331113fc1ef5e1fee20424d7d1 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 17 Jul 2024 11:05:09 +0200
Subject: [PATCH 4/8] typo

---
 core/src/allocation_optim.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 3bd348c78..2305fc317 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -861,7 +861,7 @@ function allocate_to_users_from_connected_basin!(
     (; problem) = allocation_model
     (; graph, user_demand) = p
 
-    # Get al UserDemand nodes from this subnetwork
+    # Get all UserDemand nodes from this subnetwork
     node_ids_user_demand = only(problem[:source_user].axes)
     for node_id in node_ids_user_demand
 

From d2723ad82ae8cca9143b9f271ec54a885d9a46b3 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 17 Jul 2024 11:25:01 +0200
Subject: [PATCH 5/8] Update docs

---
 docs/concept/allocation.qmd         | 2 ++
 docs/index.qmd                      | 2 +-
 docs/reference/node/user-demand.qmd | 4 ++--
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/docs/concept/allocation.qmd b/docs/concept/allocation.qmd
index c3274e15f..1585eb09e 100644
--- a/docs/concept/allocation.qmd
+++ b/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 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.
 :::
diff --git a/docs/index.qmd b/docs/index.qmd
index 1fc48c7ca..fc194805f 100644
--- a/docs/index.qmd
+++ b/docs/index.qmd
@@ -13,7 +13,7 @@ on top of the [SciML: Open Source Software for Scientific Machine Learning](http
 libraries.
 
 </br>
-Ribasim is developed and supported by Deltares.  
+Ribasim is developed and supported by Deltares.
 The first version of Ribasim is developed by Deltares as part of a consortium with funding from TKI by Deltares for the Dutch watersystem.
 </br>
 <a href="https://www.deltares.nl/">
diff --git a/docs/reference/node/user-demand.qmd b/docs/reference/node/user-demand.qmd
index 64e50ef94..8bc3913b0 100644
--- a/docs/reference/node/user-demand.qmd
+++ b/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.
 

From 483cf8ae68f197beb6de86e93329dc6ae207d29b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 17 Jul 2024 12:00:29 +0200
Subject: [PATCH 6/8] Add test

---
 core/src/allocation_optim.jl | 10 ++--------
 core/test/allocation_test.jl | 19 +++++++++++++++++++
 docs/concept/allocation.qmd  |  2 +-
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 2305fc317..ecc1bec53 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -854,11 +854,10 @@ end
 
 function allocate_to_users_from_connected_basin!(
     allocation_model::AllocationModel,
-    flow_priority,
     p::Parameters,
     priority_idx::Int,
 )::Nothing
-    (; problem) = allocation_model
+    (; flow_priority, problem) = allocation_model
     (; graph, user_demand) = p
 
     # Get all UserDemand nodes from this subnetwork
@@ -920,12 +919,7 @@ function optimize_priority!(
 
     # Allocate to UserDemand nodes from the directly connected basin
     # This happens outside the JuMP optimization
-    allocate_to_users_from_connected_basin!(
-        allocation_model,
-        flow_priority,
-        p,
-        priority_idx,
-    )
+    allocate_to_users_from_connected_basin!(allocation_model, p, priority_idx)
 
     # Solve the allocation problem for this priority
     JuMP.optimize!(problem)
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index fdde8a924..321da335a 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -550,3 +550,22 @@ end
     @test all(isapprox.(fractions[1], fractions[3], atol = 1e-4))
     @test all(isapprox.(fractions[1], fractions[4], atol = 1e-4))
 end
+
+@testitem "direct_basin_allocation" begin
+    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)
+    allocation_model.flow_priority
+    @test collect(values(res.data)) == [0.0015, 0.0, 0.0]
+end
diff --git a/docs/concept/allocation.qmd b/docs/concept/allocation.qmd
index 1585eb09e..694479724 100644
--- a/docs/concept/allocation.qmd
+++ b/docs/concept/allocation.qmd
@@ -7,7 +7,7 @@ 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 there is a simple step that tries to allocate flow to the `UserDemand` nodes from the their directly connected basin.
+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.

From d61c1445a91f559c449c823458e7d5c160e43d29 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 17 Jul 2024 13:23:48 +0200
Subject: [PATCH 7/8] Test fix

---
 core/test/allocation_test.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 321da335a..f197cc57d 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -566,6 +566,5 @@ end
     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)
-    allocation_model.flow_priority
-    @test collect(values(res.data)) == [0.0015, 0.0, 0.0]
+    @test collect(values(allocation_model.flow_priority.data)) == [0.0015, 0.0, 0.0]
 end

From cd9c351f4dc3fa2fbfe89dedc5dd41a723416c50 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 17 Jul 2024 13:41:31 +0200
Subject: [PATCH 8/8] Another test fix

---
 core/test/allocation_test.jl | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index f197cc57d..de1646e26 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -552,6 +552,7 @@ end
 end
 
 @testitem "direct_basin_allocation" begin
+    using Ribasim: NodeID
     import SQLite
     import JuMP
 
@@ -566,5 +567,8 @@ end
     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)
-    @test collect(values(allocation_model.flow_priority.data)) == [0.0015, 0.0, 0.0]
+    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