Allocation refactoring

core/src/allocation_optim.jl

@@ -524,13 +524,14 @@ Set the demand of the flow demand nodes. 2 cases:
 - Before an allocation solve, subtract the flow trough the node with a flow demand
   from the total flow demand (which will be used at the priority of the flow demand only).
 """
-function adjust_demands_user!(
+function adjust_demands!(
     allocation_model::AllocationModel,
     p::Parameters,
     priority_idx::Int,
+    user_demand::UserDemand,
 )::Nothing
     (; problem, subnetwork_id) = allocation_model
-    (; graph, user_demand) = p
+    (; graph) = p
     (; node_id, demand_reduced) = user_demand
     F = problem[:F]
 
@@ -551,7 +552,13 @@ end
 Subtract the allocated flow to the basin from its demand,
 to obtain the reduced demand used for goal programming
 """
-function adjust_demands_level!(allocation_model::AllocationModel, p::Parameters)::Nothing
+
+function adjust_demands!(
+    allocation_model::AllocationModel,
+    p::Parameters,
+    ::Int,
+    ::LevelDemand,
+)::Nothing
     (; graph, basin) = p
     (; node_id, demand) = basin
     (; subnetwork_id, problem) = allocation_model
@@ -593,8 +600,13 @@ end
 Reduce the flow demand based on flow trough the node with the demand.
 Flow from any priority counts.
 """
-function adjust_demands_flow!(allocation_model::AllocationModel, p::Parameters)::Nothing
-    (; flow_demand, graph) = p
+function adjust_demands!(
+    allocation_model::AllocationModel,
+    p::Parameters,
+    ::Int,
+    flow_demand::FlowDemand,
+)::Nothing
+    (; graph) = p
     (; problem, subnetwork_id) = allocation_model
     F = problem[:F]
 
@@ -861,7 +873,7 @@ function save_allocation_flows!(
     return nothing
 end
 
-function allocate_priority!(
+function optimize_priority!(
     allocation_model::AllocationModel,
     u::ComponentVector,
     p::Parameters,
@@ -916,9 +928,12 @@ function allocate_priority!(
     adjust_capacities_returnflow!(allocation_model, p)
 
     # Adjust demands for next optimization (in case of internal_sources -> collect_demands)
-    adjust_demands_user!(allocation_model, p, priority_idx)
-    adjust_demands_level!(allocation_model, p)
-    adjust_demands_flow!(allocation_model, p)
+    for parameter in propertynames(p)
+        demand_node = getfield(p, parameter)
+        if demand_node isa AbstractDemandNode
+            adjust_demands!(allocation_model, p, priority_idx, demand_node)
+        end
+    end
     return nothing
 end
 
@@ -968,43 +983,72 @@ 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.
 """
-function allocate!(
+function collect_demands(
     p::Parameters,
     allocation_model::AllocationModel,
     t::Float64,
     u::ComponentVector,
-    optimization_type::OptimizationType.T,
 )::Nothing
     (; allocation) = p
     (; subnetwork_id) = allocation_model
     (; priorities, subnetwork_demands) = allocation
-    main_network_source_edges = get_main_network_connections(p, subnetwork_id)
 
-    if subnetwork_id == 1
-        @assert optimization_type == OptimizationType.allocate "For the main network no demands have to be collected"
+    ## Find internal sources
+    optimization_type = OptimizationType.internal_sources
+    set_initial_capacities_inlet!(allocation_model, p, optimization_type)
+
+    set_initial_values!(allocation_model, p, u, t)
+
+    # Loop over priorities
+    for priority_idx in eachindex(priorities)
+        optimize_priority!(allocation_model, u, p, t, priority_idx, optimization_type)
     end
 
+    ## Collect demand
+    optimization_type = OptimizationType.collect_demands
+
+    main_network_source_edges = get_main_network_connections(p, subnetwork_id)
+
     # Reset the subnetwork demands to 0.0
-    if optimization_type == OptimizationType.collect_demands
-        for main_network_connection in keys(subnetwork_demands)
-            if main_network_connection in main_network_source_edges
-                subnetwork_demands[main_network_connection] .= 0.0
-            end
+    for main_network_connection in keys(subnetwork_demands)
+        if main_network_connection in main_network_source_edges
+            subnetwork_demands[main_network_connection] .= 0.0
         end
     end
 
     set_initial_capacities_inlet!(allocation_model, p, optimization_type)
 
-    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)
+    # When collecting demands, only flow should be available
+    # from the main to subnetwork connections
+    empty_sources!(allocation_model, allocation)
+
+    # Loop over priorities
+    for priority_idx in eachindex(priorities)
+        optimize_priority!(allocation_model, u, p, t, priority_idx, optimization_type)
     end
+end
+
+function allocate_demands!(
+    p::Parameters,
+    allocation_model::AllocationModel,
+    t::Float64,
+    u::ComponentVector,
+)::Nothing
+    optimization_type = OptimizationType.allocate
+    (; allocation) = p
+    (; subnetwork_id) = allocation_model
+    (; priorities) = allocation
+
+    if is_main_network(subnetwork_id)
+        @assert optimization_type == OptimizationType.allocate "For the main network no demands have to be collected"
+    end
+
+    set_initial_capacities_inlet!(allocation_model, p, optimization_type)
+
+    set_initial_values!(allocation_model, p, u, t)
 
     # Loop over the priorities
     for priority_idx in eachindex(priorities)
-        allocate_priority!(allocation_model, u, p, t, priority_idx, optimization_type)
+        optimize_priority!(allocation_model, u, p, t, priority_idx, optimization_type)
     end
 end

core/src/callback.jl

@@ -480,16 +480,15 @@ function update_allocation!(integrator)::Nothing
     # If a main network is present, collect demands of subnetworks
     if has_main_network(allocation)
         for allocation_model in Iterators.drop(allocation_models, 1)
-            allocate!(p, allocation_model, t, u, OptimizationType.internal_sources)
-            allocate!(p, allocation_model, t, u, OptimizationType.collect_demands)
+            collect_demands(p, allocation_model, t, u)
         end
     end
 
     # Solve the allocation problems
     # If a main network is present this is solved first,
     # which provides allocation to the subnetworks
     for allocation_model in allocation_models
-        allocate!(p, allocation_model, t, u, OptimizationType.allocate)
+        allocate_demands!(p, allocation_model, t, u)
     end
 
     # Reset the mean source flows

core/src/parameter.jl

@@ -135,6 +135,8 @@ end
 
 abstract type AbstractParameterNode end
 
+abstract type AbstractDemandNode <: AbstractParameterNode end
+
 """
 In-memory storage of saved mean flows for writing to results.
 
@@ -520,7 +522,7 @@ allocated: water flux currently allocated to UserDemand per priority (node_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
 """
-struct UserDemand <: AbstractParameterNode
+struct UserDemand <: AbstractDemandNode
     node_id::Vector{NodeID}
     inflow_edge::Vector{EdgeMetadata}
     outflow_edge::Vector{EdgeMetadata}
@@ -576,14 +578,14 @@ min_level: The minimum target level of the connected basin(s)
 max_level: The maximum target level of the connected basin(s)
 priority: If in a shortage state, the priority of the demand of the connected basin(s)
 """
-struct LevelDemand <: AbstractParameterNode
+struct LevelDemand <: AbstractDemandNode
     node_id::Vector{NodeID}
     min_level::Vector{ScalarInterpolation}
     max_level::Vector{ScalarInterpolation}
     priority::Vector{Int32}
 end
 
-struct FlowDemand <: AbstractParameterNode
+struct FlowDemand <: AbstractDemandNode
     node_id::Vector{NodeID}
     demand_itp::Vector{ScalarInterpolation}
     demand::Vector{Float64}

core/test/allocation_test.jl

@@ -17,7 +17,7 @@
     allocation.mean_flows[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))][] = 4.5
     allocation_model = p.allocation.allocation_models[1]
     u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
-    Ribasim.allocate!(p, allocation_model, 0.0, u, OptimizationType.allocate)
+    Ribasim.allocate_demands!(p, allocation_model, 0.0, u)
 
     # Last priority (= 2) flows
     F = allocation_model.problem[:F]
@@ -223,8 +223,7 @@ end
     # Collecting demands
     u = ComponentVector(; storage = zeros(length(basin.node_id)))
     for allocation_model in allocation_models[2:end]
-        Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources)
-        Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands)
+        Ribasim.collect_demands(p, allocation_model, t, u)
     end
 
     # See the difference between these values here and in
@@ -238,7 +237,7 @@ end
     # Solving for the main network, containing subnetworks as UserDemands
     allocation_model = allocation_models[1]
     (; problem) = allocation_model
-    Ribasim.allocate_priority!(allocation_model, u, p, t, 1, OptimizationType.allocate)
+    Ribasim.optimize_priority!(allocation_model, u, p, t, 1, OptimizationType.allocate)
 
     # Main network objective function
     F = problem[:F]
@@ -307,8 +306,7 @@ end
     # Collecting demands
     u = ComponentVector(; storage = zeros(length(basin.node_id)))
     for allocation_model in allocation_models[2:end]
-        Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources)
-        Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands)
+        Ribasim.collect_demands(p, allocation_model, t, u)
     end
 
     # See the difference between these values here and in
@@ -453,7 +451,7 @@ end
     Ribasim.set_initial_values!(allocation_model, p, u, t)
 
     # Priority 1
-    Ribasim.allocate_priority!(
+    Ribasim.optimize_priority!(
         allocation_model,
         model.integrator.u,
         p,
@@ -468,7 +466,7 @@ end
     @test JuMP.normalized_rhs(constraint_flow_out) == Inf
 
     ## Priority 2
-    Ribasim.allocate_priority!(
+    Ribasim.optimize_priority!(
         allocation_model,
         model.integrator.u,
         p,
@@ -483,7 +481,7 @@ end
     @test JuMP.normalized_rhs(constraint_flow_out) == 0.0
 
     ## Priority 3
-    Ribasim.allocate_priority!(
+    Ribasim.optimize_priority!(
         allocation_model,
         model.integrator.u,
         p,
@@ -502,7 +500,7 @@ end
     @test JuMP.value(F[(only(level_boundary.node_id), node_id_with_flow_demand)]) == 0
 
     ## Priority 4
-    Ribasim.allocate_priority!(
+    Ribasim.optimize_priority!(
         allocation_model,
         model.integrator.u,
         p,

docs/contribute/addnode.qmd

@@ -16,6 +16,14 @@ struct NewNodeType <: AbstractParameterNode
     # Other fields
 end
 ```
+Another abstract type which subtypes from `AbstractParameterNode` is called `AbstractDemandNode`. For creating new node type used in allocation, define a struct:
+
+```julia
+struct NewNodeType <: AbstractDemandNode
+    node_id::Vector{NodeID}
+    # Other fields
+end
+```
 
 These fields do not have to correspond 1:1 with the input tables (see below). The vector with all node IDs that are of the new type in a given model is a mandatory field. Now you can:
 

docs/core/allocation.qmd

@@ -20,7 +20,7 @@ The allocation algorithm contains 3 types of optimization:
 
 The full algorithm goes through the following steps:
 
-1. Perform `internal_sources` followed by `collect_demands` for all subnetworks apart from the main network;
+1. Perform `collect_demands` for all subnetworks apart from the main network;
 2. Perform `allocate` for the main network;
 3. Perform `allocate` for the other subnetworks.