Finish refactor of handling boundary sources

core/src/allocation_init.jl

@@ -95,8 +95,8 @@ function get_subnetwork_capacity(
     return capacity
 end
 
-const allocation_source_nodetypes =
-    Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary, NodeType.UserDemand])
+const boundary_source_nodetypes =
+    Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary])
 
 """
 Add the edges connecting the main network work to a subnetwork to both the main network
@@ -275,30 +275,39 @@ The constraint indices are (edge_source_id, edge_dst_id).
 Constraint:
 flow over source edge <= source flow in subnetwork
 """
-function add_constraints_source!(
+function add_constraints_boundary_source!(
     problem::JuMP.Model,
     p::Parameters,
     subnetwork_id::Int32,
 )::Nothing
-    (; graph) = p
-    edges_source = Tuple{NodeID, NodeID}[]
+    edges_source =
+        [edge for edge in source_edges_subnetwork(p, subnetwork_id) if edge[1] != edge[2]]
     F = problem[:F]
 
-    # Find the edges in the whole model which are a source for
-    # this subnetwork
-    for edge in keys(p.allocation.mean_input_flows[subnetwork_id])
-        from_id, to_id = edge
-        if graph[from_id].subnetwork_id == subnetwork_id &&
-           graph[to_id].subnetwork_id == subnetwork_id
-            push!(edges_source, edge)
-        end
-    end
+    problem[:source_boundary] = JuMP.@constraint(
+        problem,
+        [edge_id = edges_source],
+        F[edge_id] <= 0.0,
+        base_name = "source_boundary"
+    )
+    return nothing
+end
+
+function add_constraints_main_network_source!(
+    problem::JuMP.Model,
+    p::Parameters,
+    subnetwork_id::Int32,
+)::Nothing
+    F = problem[:F]
+    (; main_network_connections, subnetwork_ids) = p.allocation
+    subnetwork_id = searchsortedfirst(subnetwork_ids, subnetwork_id)
+    edges_source = main_network_connections[subnetwork_id]
 
-    problem[:source] = JuMP.@constraint(
+    problem[:source_main_network] = JuMP.@constraint(
         problem,
         [edge_id = edges_source],
         F[edge_id] <= 0.0,
-        base_name = "source"
+        base_name = "source_main_network"
     )
     return nothing
 end
@@ -448,9 +457,9 @@ function allocation_problem(
 
     # Add constraints to problem
     add_constraints_conservation_node!(problem, p, subnetwork_id)
-
     add_constraints_capacity!(problem, capacity, p, subnetwork_id)
-    add_constraints_source!(problem, p, subnetwork_id)
+    add_constraints_boundary_source!(problem, p, subnetwork_id)
+    add_constraints_main_network_source!(problem, p, subnetwork_id)
     add_constraints_user_source!(problem, p, subnetwork_id)
     add_constraints_basin_flow!(problem)
     add_constraints_buffer!(problem)
@@ -481,12 +490,12 @@ function get_sources_in_order(
         sources[edge] = AllocationSource(; edge, type = AllocationSourceType.user_return)
     end
 
-    # Source edges (within subnetwork)
-    for edge in
-        sort(only(problem[:source].axes); by = edge -> (edge[1].value, edge[2].value))
-        if graph[edge[1]].subnetwork_id == graph[edge[2]].subnetwork_id
-            sources[edge] = AllocationSource(; edge, type = AllocationSourceType.edge)
-        end
+    # Boundary node sources
+    for edge in sort(
+        only(problem[:source_boundary].axes);
+        by = edge -> (edge[1].value, edge[2].value),
+    )
+        sources[edge] = AllocationSource(; edge, type = AllocationSourceType.boundary_node)
     end
 
     # Basins with level demand

core/src/allocation_optim.jl

@@ -127,8 +127,7 @@ function assign_allocations!(
         # If this edge is a source edge from the main network to a subnetwork,
         # and demands are being collected, add its flow to the demand of this edge
         if optimization_type == OptimizationType.collect_demands
-            if graph[edge...].subnetwork_id_source == subnetwork_id &&
-               edge ∈ main_network_source_edges
+            if edge in main_network_source_edges
                 allocated = flow[edge]
                 subnetwork_demands[edge][priority_idx] += allocated
             end
@@ -206,16 +205,12 @@ function set_initial_capacities_source!(
     (; allocation) = p
     (; mean_input_flows) = allocation
     (; subnetwork_id) = allocation_model
-    main_network_source_edges = get_main_network_connections(p, subnetwork_id)
 
-    for edge in keys(p.allocation.mean_input_flows[subnetwork_id])
-        # If it is a source edge for this allocation problem
-        if edge ∉ main_network_source_edges
-            # Reset the source to the averaged flow over the last allocation period
-            source = sources[edge]
-            @assert source.type == AllocationSourceType.edge
-            source.capacity = mean_input_flows[subnetwork_id][edge]
-        end
+    mean_input_flows_subnetwork_ = mean_input_flows_subnetwork(p, subnetwork_id)
+
+    for edge in keys(mean_input_flows_subnetwork_)
+        source = sources[edge]
+        source.capacity = mean_input_flows_subnetwork_[edge]
     end
     return nothing
 end
@@ -228,7 +223,7 @@ function reduce_source_capacity!(problem::JuMP.Model, source::AllocationSource):
 
     used_capacity =
         if source.type in (
-            AllocationSourceType.edge,
+            AllocationSourceType.boundary_node,
             AllocationSourceType.main_to_sub,
             AllocationSourceType.user_return,
         )
@@ -340,11 +335,10 @@ function get_basin_data(
     u::ComponentVector,
     node_id::NodeID,
 )
-    (; graph, allocation, basin) = p
-    (; Δt_allocation) = allocation_model
-    (; mean_input_flows) = allocation
+    (; graph, basin) = p
+    (; Δt_allocation, subnetwork_id) = allocation_model
     @assert node_id.type == NodeType.Basin
-    influx = mean_input_flows[(node_id, node_id)][]
+    influx = mean_input_flows_subnetwork(p, subnetwork_id)[(node_id, node_id)]
     storage_basin = basin.current_properties.current_storage[parent(u)][node_id.idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)
     if isempty(control_inneighbors)
@@ -840,9 +834,10 @@ function set_source_capacity!(
     optimization_type::OptimizationType.T,
 )::Nothing
     (; problem, sources) = allocation_model
-    constraints_source_edge = problem[:source]
-    constraints_source_basin = problem[:basin_outflow]
+    constraints_source_boundary = problem[:source_boundary]
     constraints_source_user_out = problem[:source_user]
+    constraints_source_main_network = problem[:source_main_network]
+    constraints_source_basin = problem[:basin_outflow]
     constraints_source_buffer = problem[:flow_buffer_outflow]
 
     for source in values(sources)
@@ -859,16 +854,17 @@ function set_source_capacity!(
             0.0
         end
 
-        constraint =
-            if source.type in (AllocationSourceType.edge, AllocationSourceType.main_to_sub)
-                constraints_source_edge[edge]
-            elseif source.type == AllocationSourceType.basin
-                constraints_source_basin[edge[1]]
-            elseif source.type == AllocationSourceType.user_return
-                constraints_source_user_out[edge[1]]
-            elseif source.type == AllocationSourceType.buffer
-                constraints_source_buffer[edge[1]]
-            end
+        constraint = if source.type == AllocationSourceType.boundary_node
+            constraints_source_boundary[edge]
+        elseif source.type == AllocationSourceType.main_to_sub
+            constraints_source_main_network[edge]
+        elseif source.type == AllocationSourceType.basin
+            constraints_source_basin[edge[1]]
+        elseif source.type == AllocationSourceType.user_return
+            constraints_source_user_out[edge[1]]
+        elseif source.type == AllocationSourceType.buffer
+            constraints_source_buffer[edge[1]]
+        end
 
         JuMP.set_normalized_rhs(constraint, capacity_effective)
     end
@@ -1043,7 +1039,8 @@ function empty_sources!(allocation_model::AllocationModel, allocation::Allocatio
     (; problem) = allocation_model
     (; subnetwork_demands) = allocation
 
-    for constraint_set_name in [:source, :source_user, :basin_outflow, :flow_buffer_outflow]
+    for constraint_set_name in
+        [:source_boundary, :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

core/src/callback.jl

@@ -153,8 +153,11 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
     end
 
     # Update realized flows for allocation input
-    for edge in keys(allocation.mean_input_flows)
-        allocation.mean_input_flows[edge] += flow_update_on_edge(integrator, edge)
+    for subnetwork_id in allocation.subnetwork_ids
+        mean_input_flows_subnetwork_ = mean_input_flows_subnetwork(p, subnetwork_id)
+        for edge in keys(mean_input_flows_subnetwork_)
+            mean_input_flows_subnetwork_[edge] += flow_update_on_edge(integrator, edge)
+        end
     end
 
     # Update realized flows for allocation output
@@ -773,8 +776,10 @@ function update_allocation!(integrator)::Nothing
 
     # Divide by the allocation Δt to get the mean input flows from the cumulative flows
     (; Δt_allocation) = allocation_models[1]
-    for edge in keys(mean_input_flows)
-        mean_input_flows[edge] /= Δt_allocation
+    for mean_input_flows_subnetwork in values(mean_input_flows)
+        for edge in keys(mean_input_flows_subnetwork)
+            mean_input_flows_subnetwork[edge] /= Δt_allocation
+        end
     end
 
     # Divide by the allocation Δt to get the mean realized flows from the cumulative flows
@@ -797,7 +802,7 @@ function update_allocation!(integrator)::Nothing
     end
 
     # Reset the mean flows
-    for mean_flows in (mean_input_flows, mean_realized_flows)
+    for mean_flows in (mean_input_flows..., mean_realized_flows)
         for edge in keys(mean_flows)
             mean_flows[edge] = 0.0
         end

core/src/parameter.jl

@@ -184,13 +184,14 @@ record_flow: A record of all flows computed by allocation optimization, eventual
     output file
 """
 @kwdef struct Allocation
-    subnetwork_ids::Vector{Int32} = []
-    allocation_models::Vector{AllocationModel} = []
-    main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}} = []
+    subnetwork_ids::Vector{Int32} = Int32[]
+    allocation_models::Vector{AllocationModel} = AllocationModel[]
+    main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}} =
+        Vector{Tuple{NodeID, NodeID}}[]
     priorities::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()
-    mean_input_flows::Dict{Int32, Dict{Tuple{NodeID, NodeID}, Float64}}
+    mean_input_flows::Vector{Dict{Tuple{NodeID, NodeID}, Float64}}
     mean_realized_flows::Dict{Tuple{NodeID, NodeID}, Float64}
     record_demand::@NamedTuple{
         time::Vector{Float64},

core/src/read.jl

@@ -1268,25 +1268,31 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
 end
 
 function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
-    mean_input_flows = Dict{Int32, Dict{Tuple{NodeID, NodeID}, Float64}}()
+    mean_input_flows = Dict{Tuple{NodeID, NodeID}, Float64}[]
+
+    subnetwork_ids = sort(collect(keys(graph[].node_ids)))
+
+    for subnetwork_id in subnetwork_ids
+        push!(mean_input_flows, Dict{Tuple{NodeID, NodeID}, Float64}())
+    end
 
     # Find edges which serve as sources in allocation
     for edge_metadata in values(graph.edge_data)
         (; edge) = edge_metadata
         id_source, _ = edge
-        if id_source in allocation_source_nodetypes
+        if id_source.type in boundary_source_nodetypes
             (; subnetwork_id) = graph[id_source]
-            if !haskey(mean_input_flows, subnetwork_id)
-                mean_input_flows[subnetwork_id] = Dict{Tuple{NodeID, NodeID}, Float64}()
-            end
-            mean_input_flows[subnetwork_id][edge] = 0.0
+            subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+            mean_input_flows[subnetwork_idx][edge] = 0.0
         end
     end
 
     # Find basins with a level demand
     for node_id in values(graph.vertex_labels)
         if has_external_demand(graph, node_id, :level_demand)[1]
-            mean_input_flows[(node_id, node_id)] = 0.0
+            subnetwork_id = graph[node_id].subnetwork_id
+            subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+            mean_input_flows[subnetwork_idx][(node_id, node_id)] = 0.0
         end
     end
 

core/src/util.jl

@@ -587,15 +587,17 @@ function set_initial_allocation_mean_flows!(integrator)::Nothing
     du = get_du(integrator)
     water_balance!(du, u, p, t)
 
-    for edge in keys(mean_input_flows)
-        if edge[1] == edge[2]
-            q = get_influx(du, edge[1], p)
-        else
-            q = get_flow(du, p, t, edge)
+    for mean_input_flows_subnetwork in values(mean_input_flows)
+        for edge in keys(mean_input_flows_subnetwork)
+            if edge[1] == edge[2]
+                q = get_influx(du, edge[1], p)
+            else
+                q = get_flow(du, p, t, edge)
+            end
+            # Multiply by Δt_allocation as averaging divides by this factor
+            # in update_allocation!
+            mean_input_flows_subnetwork[edge] = q * Δt_allocation
         end
-        # Multiply by Δt_allocation as averaging divides by this factor
-        # in update_allocation!
-        mean_input_flows[edge] = q * Δt_allocation
     end
 
     # Mean realized demands for basins are calculated as Δstorage/Δt
@@ -1168,3 +1170,12 @@ function min_low_user_demand_level_factor(
         one(T)
     end
 end
+
+function mean_input_flows_subnetwork(p::Parameters, subnetwork_id::Int32)
+    (; mean_input_flows, subnetwork_ids) = p.allocation
+    subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+    return mean_input_flows[subnetwork_idx]
+end
+
+source_edges_subnetwork(p::Parameters, subnetwork_id::Int32) =
+    keys(mean_input_flows_subnetwork(p, subnetwork_id))