Use mean input in allocation (#1427)

Fixes #1121.

To do:
- [x] Use average source flows
- [x] Use average vertical fluxes on basin for level demand
- [x] Update documentation

core/src/allocation_init.jl

@@ -153,7 +153,8 @@ function add_variables_basin!(
     (; graph) = p
     node_ids_basin = [
         node_id for
-        node_id in graph[].node_ids[subnetwork_id] if graph[node_id].type == :basin
+        node_id in graph[].node_ids[subnetwork_id] if graph[node_id].type == :basin &&
+        has_external_demand(graph, node_id, :level_demand)[1]
     ]
     problem[:F_basin_in] =
         JuMP.@variable(problem, F_basin_in[node_id = node_ids_basin,] >= 0.0)
@@ -211,7 +212,7 @@ function add_variables_absolute_value!(
         type = node_id.type
         if type == NodeType.UserDemand
             push!(node_ids_user_demand, node_id)
-        elseif type == NodeType.Basin
+        elseif has_external_demand(graph, node_id, :level_demand)[1]
             push!(node_ids_level_demand, node_id)
         elseif has_external_demand(graph, node_id, :flow_demand)[1]
             push!(node_ids_flow_demand, node_id)
@@ -392,8 +393,8 @@ function add_constraints_conservation_node!(
             end
         end
 
-        # If the node is a Basin, add basin in- and outflow
-        if node_id.type == NodeType.Basin
+        # If the node is a Basin with a level demand, add basin in- and outflow
+        if has_external_demand(graph, node_id, :level_demand)[1]
             push!(inflows_node, F_basin_out[node_id])
             push!(outflows_node, F_basin_in[node_id])
         end

core/src/allocation_optim.jl

@@ -254,7 +254,8 @@ function set_initial_capacities_source!(
     p::Parameters,
 )::Nothing
     (; problem) = allocation_model
-    (; graph) = p
+    (; graph, allocation) = p
+    (; mean_flows) = allocation
     (; subnetwork_id) = allocation_model
     source_constraints = problem[:source]
     main_network_source_edges = get_main_network_connections(p, subnetwork_id)
@@ -264,8 +265,8 @@ function set_initial_capacities_source!(
         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 current flow from the physical layer.
-                source_capacity = get_flow(graph, edge..., 0)
+                # Reset the source to the averaged flow over the last allocation period
+                source_capacity = mean_flows[edge][]
                 JuMP.set_normalized_rhs(
                     source_constraints[edge],
                     # It is assumed that the allocation procedure does not have to be differentiated.
@@ -357,14 +358,14 @@ function get_basin_data(
     u::ComponentVector,
     node_id::NodeID,
 )
-    (; graph, basin, level_demand) = p
+    (; graph, basin, level_demand, allocation) = p
     (; vertical_flux) = basin
     (; Δt_allocation) = allocation_model
+    (; mean_flows) = allocation
     @assert node_id.type == NodeType.Basin
     vertical_flux = get_tmp(vertical_flux, 0)
     _, basin_idx = id_index(basin.node_id, node_id)
-    # NOTE: Instantaneous
-    influx = get_influx(basin, node_id)
+    influx = mean_flows[(node_id, node_id)][]
     _, basin_idx = id_index(basin.node_id, node_id)
     storage_basin = u.storage[basin_idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)
@@ -511,12 +512,15 @@ function set_initial_demands_level!(
     p::Parameters,
     t::Float64,
 )::Nothing
-    (; subnetwork_id) = allocation_model
+    (; subnetwork_id, problem) = allocation_model
     (; graph, basin) = p
     (; node_id, demand) = basin
 
-    for (i, id) in enumerate(node_id)
+    node_ids_level_demand = only(problem[:basin_outflow].axes)
+
+    for id in node_ids_level_demand
         if graph[id].subnetwork_id == subnetwork_id
+            _, i = id_index(node_id, id)
             demand[i] = get_basin_demand(allocation_model, u, p, t, id)
         end
     end
@@ -606,8 +610,9 @@ function adjust_demands_level!(allocation_model::AllocationModel, p::Parameters)
     F_basin_in = problem[:F_basin_in]
 
     # Reduce the demand by what was allocated
-    for (i, id) in enumerate(node_id)
+    for id in only(F_basin_in.axes)
         if graph[id].subnetwork_id == subnetwork_id
+            _, i = id_index(basin.node_id, id)
             demand[i] -= JuMP.value(F_basin_in[id])
         end
     end
@@ -763,8 +768,7 @@ function save_demands_and_allocations!(
             #NOTE: instantaneous
             realized = get_flow(graph, inflow_id(graph, node_id), node_id, 0)
 
-        elseif node_id.type == NodeType.Basin &&
-               has_external_demand(graph, node_id, :level_demand)[1]
+        elseif has_external_demand(graph, node_id, :level_demand)[1]
             basin_priority_idx = get_external_priority_idx(p, node_id)
 
             if priority_idx == 1 || basin_priority_idx == priority_idx

core/src/callback.jl

@@ -91,7 +91,7 @@ Integrate flows over the last timestep
 """
 function integrate_flows!(u, t, integrator)::Nothing
     (; p, dt) = integrator
-    (; graph, user_demand, basin) = p
+    (; graph, user_demand, basin, allocation) = p
     (; flow, flow_dict, flow_prev, flow_integrated) = graph[]
     (; vertical_flux, vertical_flux_prev, vertical_flux_integrated, vertical_flux_bmi) =
         basin
@@ -106,12 +106,31 @@ function integrate_flows!(u, t, integrator)::Nothing
     @. vertical_flux_integrated += 0.5 * (vertical_flux + vertical_flux_prev) * dt
     @. vertical_flux_bmi += 0.5 * (vertical_flux + vertical_flux_prev) * dt
 
+    # UserDemand realized flows for BMI
     for (i, id) in enumerate(user_demand.node_id)
         src_id = inflow_id(graph, id)
         flow_idx = flow_dict[src_id, id]
         user_demand.realized_bmi[i] += 0.5 * (flow[flow_idx] + flow_prev[flow_idx]) * dt
     end
 
+    # Allocation source flows
+    for (edge, value) in allocation.mean_flows
+        if edge[1] == edge[2]
+            # Vertical fluxes
+            _, basin_idx = id_index(basin.node_id, edge[1])
+            value[] +=
+                0.5 *
+                (get_influx(basin, basin_idx) + get_influx(basin, basin_idx; prev = true)) *
+                dt
+        else
+            # Horizontal flows
+            value[] +=
+                0.5 *
+                (get_flow(graph, edge..., 0) + get_flow(graph, edge..., 0; prev = true)) *
+                dt
+        end
+    end
+
     copyto!(flow_prev, flow)
     copyto!(vertical_flux_prev, vertical_flux)
     return nothing
@@ -446,7 +465,14 @@ end
 function update_allocation!(integrator)::Nothing
     (; p, t, u) = integrator
     (; allocation) = p
-    (; allocation_models) = allocation
+    (; allocation_models, mean_flows) = allocation
+    (; Δt_allocation) = allocation_models[1]
+
+    # Divide by the allocation Δt to obtain the mean flows
+    # from the integrated flows
+    for value in values(mean_flows)
+        value[] /= Δt_allocation
+    end
 
     # If a main network is present, collect demands of subnetworks
     if has_main_network(allocation)
@@ -462,6 +488,11 @@ function update_allocation!(integrator)::Nothing
     for allocation_model in allocation_models
         allocate!(p, allocation_model, t, u, OptimizationType.allocate)
     end
+
+    # Reset the mean source flows
+    for value in values(mean_flows)
+        value[] = 0.0
+    end
 end
 
 "Load updates from 'TabulatedRatingCurve / time' into the parameters"

core/src/graph.jl

@@ -33,16 +33,15 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         node_id = NodeID(row.node_type, row.node_id)
         # Process allocation network ID
         if ismissing(row.subnetwork_id)
-            allocation_network_id = 0
+            subnetwork_id = 0
         else
-            allocation_network_id = row.subnetwork_id
-            if !haskey(node_ids, allocation_network_id)
-                node_ids[allocation_network_id] = Set{NodeID}()
+            subnetwork_id = row.subnetwork_id
+            if !haskey(node_ids, subnetwork_id)
+                node_ids[subnetwork_id] = Set{NodeID}()
             end
-            push!(node_ids[allocation_network_id], node_id)
+            push!(node_ids[subnetwork_id], node_id)
         end
-        graph[node_id] =
-            NodeMetadata(Symbol(snake_case(row.node_type)), allocation_network_id)
+        graph[node_id] = NodeMetadata(Symbol(snake_case(row.node_type)), subnetwork_id)
     end
 
     errors = false
@@ -178,9 +177,16 @@ end
 """
 Get the flow over the given edge (val is needed for get_tmp from ForwardDiff.jl).
 """
-function get_flow(graph::MetaGraph, id_src::NodeID, id_dst::NodeID, val)::Number
-    (; flow_dict, flow) = graph[]
-    return get_tmp(flow, val)[flow_dict[id_src, id_dst]]
+function get_flow(
+    graph::MetaGraph,
+    id_src::NodeID,
+    id_dst::NodeID,
+    val;
+    prev::Bool = false,
+)::Number
+    (; flow_dict, flow, flow_prev) = graph[]
+    flow_vector = prev ? flow_prev : flow
+    return get_tmp(flow_vector, val)[flow_dict[id_src, id_dst]]
 end
 
 """

core/src/parameter.jl

@@ -43,7 +43,7 @@ const VectorInterpolation =
 """
 Store information for a subnetwork used for allocation.
 
-allocation_network_id: The ID of this allocation network
+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
 problem: The JuMP.jl model for solving the allocation problem
 Δt_allocation: The time interval between consecutive allocation solves
@@ -57,13 +57,15 @@ end
 
 """
 Object for all information about allocation
-allocation_network_ids: The unique sorted allocation network IDs
+subnetwork_ids: The unique sorted allocation network IDs
 allocation models: The allocation models for the main network and subnetworks corresponding to
-    allocation_network_ids
+    subnetwork_ids
 main_network_connections: (from_id, to_id) from the main network to the subnetwork per subnetwork
 priorities: All used priority values.
 subnetwork_demands: The demand of an edge from the main network to a subnetwork
-record_demand: A record of demands and allocated flows for nodes that have these.
+subnetwork_allocateds: The allocated flow of an edge from the main network to a subnetwork
+mean_flows: Flows averaged over Δt_allocation over edges that are allocation sources
+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
 """
@@ -74,6 +76,7 @@ struct Allocation
     priorities::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
+    mean_flows::Dict{Tuple{NodeID, NodeID}, Base.RefValue{Float64}}
     record_demand::@NamedTuple{
         time::Vector{Float64},
         subnetwork_id::Vector{Int32},
@@ -103,7 +106,7 @@ is_active(allocation::Allocation) = !isempty(allocation.allocation_models)
 """
 Type for storing metadata of nodes in the graph
 type: type of the node
-allocation_network_id: Allocation network ID (0 if not in subnetwork)
+subnetwork_id: Allocation network ID (0 if not in subnetwork)
 """
 struct NodeMetadata
     type::Symbol

core/src/read.jl

@@ -951,7 +951,7 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
     return Subgrid(basin_ids, interpolations, fill(NaN, length(basin_ids)))
 end
 
-function Allocation(db::DB, config::Config)::Allocation
+function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
     record_demand = (
         time = Float64[],
         subnetwork_id = Int32[],
@@ -976,18 +976,34 @@ function Allocation(db::DB, config::Config)::Allocation
         optimization_type = String[],
     )
 
-    allocation = Allocation(
+    mean_flows = Dict{Tuple{NodeID, NodeID}, Base.RefValue{Float64}}()
+
+    # Find edges which serve as sources in allocation
+    for edge_metadata in values(graph.edge_data)
+        (; subnetwork_id_source, edge) = edge_metadata
+        if subnetwork_id_source != 0
+            mean_flows[edge] = Ref(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_flows[(node_id, node_id)] = Ref(0.0)
+        end
+    end
+
+    return Allocation(
         Int32[],
         AllocationModel[],
         Vector{Tuple{NodeID, NodeID}}[],
         get_all_priorities(db, config),
-        Dict{Tuple{NodeID, NodeID}, Float64}(),
-        Dict{Tuple{NodeID, NodeID}, Float64}(),
+        Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
+        Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
+        mean_flows,
         record_demand,
         record_flow,
     )
-
-    return allocation
 end
 
 """
@@ -1007,7 +1023,7 @@ function Parameters(db::DB, config::Config)::Parameters
     n_states = length(get_ids(db, "Basin")) + length(get_ids(db, "PidControl"))
     chunk_sizes = get_chunk_sizes(config, n_states)
     graph = create_graph(db, config, chunk_sizes)
-    allocation = Allocation(db, config)
+    allocation = Allocation(db, config, graph)
 
     if !valid_edges(graph)
         error("Invalid edge(s) found.")

core/src/util.jl

@@ -711,10 +711,11 @@ function get_influx(basin::Basin, node_id::NodeID)::Float64
     return get_influx(basin, basin_idx)
 end
 
-function get_influx(basin::Basin, basin_idx::Int)::Float64
-    (; vertical_flux) = basin
+function get_influx(basin::Basin, basin_idx::Int; prev::Bool = false)::Float64
+    (; vertical_flux, vertical_flux_prev) = basin
     vertical_flux = get_tmp(vertical_flux, 0)
-    (; precipitation, evaporation, drainage, infiltration) = vertical_flux
+    flux_vector = prev ? vertical_flux_prev : vertical_flux
+    (; precipitation, evaporation, drainage, infiltration) = flux_vector
     return precipitation[basin_idx] - evaporation[basin_idx] + drainage[basin_idx] -
            infiltration[basin_idx]
 end