Get rid of specifying sources in allocation by edges

core/src/allocation_init.jl

@@ -95,7 +95,7 @@
     return capacity
 end
 
-const allocation_source_nodetypes =
+const boundary_source_nodetypes =
     Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary])
 
 """
@@ -138,10 +138,6 @@
     capacity = get_subnetwork_capacity(p, subnetwork_id)
     add_subnetwork_connections!(capacity, p, subnetwork_id)
 
-    if !valid_sources(p, capacity, subnetwork_id)
-        error("Errors in sources in allocation network.")
-    end
-
     return capacity
 end
 
@@ -272,36 +268,55 @@
 end
 
 """
-Add the source constraints to the allocation problem.
+Add the boundary source constraints to the allocation problem.
 The actual threshold values will be set before each allocation solve.
 The constraint indices are (edge_source_id, edge_dst_id).
 
 Constraint:
-flow over source edge <= source flow in subnetwork
+flow over source edge <= source flow in physical layer
 """
-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}[]
+    # Source edges (without the basins)
+    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_metadata in values(graph.edge_data)
-        (; edge) = edge_metadata
-        if graph[edge...].subnetwork_id_source == 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
 
-    problem[:source] = JuMP.@constraint(
+"""
+Add main network source constraints to the allocation problem.
+The actual threshold values will be set before each allocation solve.
+The constraint indices are (edge_source_id, edge_dst_id).
+
+Constraint:
+flow over main network to subnetwork connection edge <= either 0 or allocated amount from the main network
+"""
+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_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
@@ -451,9 +466,9 @@
 
     # 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)
@@ -484,12 +499,12 @@
         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
@@ -203,22 +202,13 @@ function set_initial_capacities_source!(
     p::Parameters,
 )::Nothing
     (; sources) = allocation_model
-    (; graph, allocation) = p
-    (; mean_input_flows) = allocation
     (; subnetwork_id) = allocation_model
-    main_network_source_edges = get_main_network_connections(p, subnetwork_id)
 
-    for edge_metadata in values(graph.edge_data)
-        (; edge) = edge_metadata
-        if graph[edge...].subnetwork_id_source == 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[edge][]
-            end
-        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
@@ -231,7 +221,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,
         )
@@ -343,11 +333,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)
@@ -843,9 +832,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)
@@ -862,16 +852,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
@@ -1046,7 +1037,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/graph.jl

@@ -28,8 +28,7 @@ function create_graph(db::DB, config::Config)::MetaGraph
     )
     # Node IDs per subnetwork
     node_ids = Dict{Int32, Set{NodeID}}()
-    # Source edges per subnetwork
-    edges_source = Dict{Int32, Set{EdgeMetadata}}()
+
     # The metadata of the flow edges in the order in which they are in the input
     # and will be in the output
     flow_edges = EdgeMetadata[]
@@ -57,15 +56,8 @@ function create_graph(db::DB, config::Config)::MetaGraph
     end
 
     errors = false
-    for (;
-        edge_id,
-        from_node_type,
-        from_node_id,
-        to_node_type,
-        to_node_id,
-        edge_type,
-        subnetwork_id,
-    ) in edge_rows
+    for (; edge_id, from_node_type, from_node_id, to_node_type, to_node_id, edge_type) in
+        edge_rows
         try
             # hasfield does not work
             edge_type = getfield(EdgeType, Symbol(edge_type))
@@ -74,15 +66,8 @@ function create_graph(db::DB, config::Config)::MetaGraph
         end
         id_src = NodeID(from_node_type, from_node_id, db)
         id_dst = NodeID(to_node_type, to_node_id, db)
-        if ismissing(subnetwork_id)
-            subnetwork_id = 0
-        end
-        edge_metadata = EdgeMetadata(;
-            id = edge_id,
-            type = edge_type,
-            subnetwork_id_source = subnetwork_id,
-            edge = (id_src, id_dst),
-        )
+        edge_metadata =
+            EdgeMetadata(; id = edge_id, type = edge_type, edge = (id_src, id_dst))
         if edge_type == EdgeType.flow
             push!(flow_edges, edge_metadata)
         end
@@ -91,12 +76,6 @@ function create_graph(db::DB, config::Config)::MetaGraph
             @error "Duplicate edge" id_src id_dst
         end
         graph[id_src, id_dst] = edge_metadata
-        if subnetwork_id != 0
-            if !haskey(edges_source, subnetwork_id)
-                edges_source[subnetwork_id] = Set{EdgeMetadata}()
-            end
-            push!(edges_source[subnetwork_id], edge_metadata)
-        end
     end
     if errors
         error("Invalid edges found")
@@ -106,7 +85,7 @@ function create_graph(db::DB, config::Config)::MetaGraph
         error("Incomplete connectivity in subnetwork")
     end
 
-    graph_data = (; node_ids, edges_source, flow_edges, config.solver.saveat)
+    graph_data = (; node_ids, flow_edges, config.solver.saveat)
     @reset graph.graph_data = graph_data
 
     return graph

core/src/parameter.jl

@@ -128,7 +128,7 @@ The caches are always initialized with zeros
 """
 cache(len::Int)::Cache = LazyBufferCache(Returns(len); initializer! = set_zero!)
 
-@enumx AllocationSourceType edge basin main_to_sub user_return buffer
+@enumx AllocationSourceType boundary_node basin main_to_sub user_return buffer
 
 """
 Data structure for a single source within an allocation subnetwork.
@@ -177,20 +177,21 @@ main_network_connections: (from_id, to_id) from the main network to the subnetwo
 priorities: All used priority values.
 subnetwork_demands: The demand of an edge from the main network to a subnetwork
 subnetwork_allocateds: The allocated flow of an edge from the main network to a subnetwork
-mean_input_flows: Flows averaged over Δt_allocation over edges that are allocation sources
+mean_input_flows: Per subnetwork, flows averaged over Δt_allocation over edges that are allocation sources
 mean_realized_flows: Flows averaged over Δt_allocation over edges that realize a demand
 record_demand: A record of demands and allocated flows for nodes that have these
 record_flow: A record of all flows computed by allocation optimization, eventually saved to
     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{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},
@@ -252,14 +253,11 @@ end
 Type for storing metadata of edges in the graph:
 id: ID of the edge (only used for labeling flow output)
 type: type of the edge
-subnetwork_id_source: ID of subnetwork where this edge is a source
-  (0 if not a source)
 edge: (from node ID, to node ID)
 """
 @kwdef struct EdgeMetadata
     id::Int32
     type::EdgeType.T
-    subnetwork_id_source::Int32
     edge::Tuple{NodeID, NodeID}
 end
 
@@ -899,7 +897,6 @@ const ModelGraph = MetaGraph{
     EdgeMetadata,
     @NamedTuple{
         node_ids::Dict{Int32, Set{NodeID}},
-        edges_source::Dict{Int32, Set{EdgeMetadata}},
         flow_edges::Vector{EdgeMetadata},
         saveat::Float64,
     },