From 1b8f1759a115d008cb1dd925ac85bb69e4efe894 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 25 Apr 2024 16:02:41 +0200
Subject: [PATCH] Collect demands with empty sources (apart from the main
 network inlet)

---
 core/src/allocation_optim.jl | 28 ++++++++++++++++++++++++++--
 core/test/allocation_test.jl |  4 ++--
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 01370c228..5c306c6ac 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -247,7 +247,7 @@ end
 
 """
 Set the capacities of the sources in the subnetwork
-as the latest instantaneous flow out of the source in the physical layer
+as the average flow over the last Δt_allocation of the source in the physical layer
 """
 function set_initial_capacities_source!(
     allocation_model::AllocationModel,
@@ -996,6 +996,26 @@ function set_initial_values!(
     return nothing
 end
 
+"""
+Set the capacities of all edges that denote a source to 0.0.
+"""
+function empty_sources!(allocation_model::AllocationModel, allocation::Allocation)::Nothing
+    (; problem) = allocation_model
+    (; subnetwork_demands) = allocation
+
+    for constraint_set_name in [:source, :source_user, :basin_outflow, :flow_buffer_outflow]
+        constraint_set = problem[constraint_set_name]
+        for key in only(constraint_set.axes)
+            # Do not set the capacity to 0.0 if the edge
+            # is a main to subnetwork connection edge
+            if key ∉ keys(subnetwork_demands)
+                JuMP.set_normalized_rhs(constraint_set[key], 0.0)
+            end
+        end
+    end
+    return nothing
+end
+
 """
 Update the allocation optimization problem for the given subnetwork with the problem state
 and flows, solve the allocation problem and assign the results to the UserDemand.
@@ -1027,7 +1047,11 @@ function allocate!(
 
     set_initial_capacities_inlet!(allocation_model, p, optimization_type)
 
-    if optimization_type != OptimizationType.collect_demands
+    if optimization_type == OptimizationType.collect_demands
+        # When collecting demands, only flow should be available
+        # from the main to subnetwork connections
+        empty_sources!(allocation_model, allocation)
+    else
         set_initial_values!(allocation_model, p, u, t)
     end
 
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index d49f0512d..c89840d3c 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -271,7 +271,7 @@ end
     @test user_demand.allocated[7, :] ≈ [0.001, 0.0, 0.0]
 end
 
-@testitem "subnetworks with sources" begin
+@testitem "Subnetworks with sources" begin
     using SQLite
     using Ribasim: NodeID, OptimizationType
     using ComponentArrays: ComponentVector
@@ -307,7 +307,7 @@ end
     # See the difference between these values here and in
     # "allocation with main network optimization problem", internal sources
     # lower the subnetwork demands
-    @test subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))] ≈ [3.1, 4.0, 0.0]
+    @test subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))] ≈ [4.0, 4.0, 0.0]
     @test subnetwork_demands[(NodeID(:Basin, 6), NodeID(:Pump, 24))] ≈ [0.004, 0.0, 0.0]
     @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2] ≈ [0.001, 0.001]
 end