From eeba58a7f5f7b7bddb8700c82ededd10cfe3e42f Mon Sep 17 00:00:00 2001 From: Bart de Koning Date: Wed, 27 Nov 2024 12:20:47 +0100 Subject: [PATCH] Finish refactor of handling boundary sources --- core/src/allocation_init.jl | 57 +++++++++++++++++++++--------------- core/src/allocation_optim.jl | 55 ++++++++++++++++------------------ core/src/callback.jl | 15 ++++++---- core/src/parameter.jl | 9 +++--- core/src/read.jl | 20 ++++++++----- core/src/util.jl | 27 ++++++++++++----- core/test/allocation_test.jl | 57 ++++++++++++++++++------------------ 7 files changed, 135 insertions(+), 105 deletions(-) diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl index 144976a0c..8d63b8d94 100644 --- a/core/src/allocation_init.jl +++ b/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 diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl index 20f47fdfc..9b2a0eac7 100644 --- a/core/src/allocation_optim.jl +++ b/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 diff --git a/core/src/callback.jl b/core/src/callback.jl index 213e97284..cda123cff 100644 --- a/core/src/callback.jl +++ b/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 diff --git a/core/src/parameter.jl b/core/src/parameter.jl index 292be86e7..62e3f76e8 100644 --- a/core/src/parameter.jl +++ b/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}, diff --git a/core/src/read.jl b/core/src/read.jl index 0c565126d..d8e43e3eb 100644 --- a/core/src/read.jl +++ b/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 diff --git a/core/src/util.jl b/core/src/util.jl index 852595d92..2eac3f267 100644 --- a/core/src/util.jl +++ b/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)) diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl index 6da2fec36..3dd6c7ba0 100644 --- a/core/test/allocation_test.jl +++ b/core/test/allocation_test.jl @@ -14,7 +14,8 @@ (; graph, allocation) = p - allocation.mean_input_flows[(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] = 4.5 + allocation.mean_input_flows[1][(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] = + 4.5 allocation_model = p.allocation.allocation_models[1] (; flow) = allocation_model u = ComponentVector(; storage = zeros(length(p.basin.node_id))) @@ -105,11 +106,11 @@ end # In each subnetwork, the connection from the main network to the subnetwork is # interpreted as a source - @test Ribasim.get_allocation_model(p, Int32(3)).problem[:source].axes[1] == + @test Ribasim.get_allocation_model(p, Int32(3)).problem[:source_main_network].axes[1] == [(NodeID(:Basin, 2, p), NodeID(:Pump, 11, p))] - @test Ribasim.get_allocation_model(p, Int32(5)).problem[:source].axes[1] == + @test Ribasim.get_allocation_model(p, Int32(5)).problem[:source_main_network].axes[1] == [(NodeID(:Basin, 6, p), NodeID(:Pump, 24, p))] - @test Ribasim.get_allocation_model(p, Int32(7)).problem[:source].axes[1] == + @test Ribasim.get_allocation_model(p, Int32(7)).problem[:source_main_network].axes[1] == [(NodeID(:Basin, 10, p), NodeID(:Pump, 38, p))] end @@ -174,7 +175,7 @@ end # Running full allocation algorithm (; Δt_allocation) = allocation_models[1] - mean_input_flows[(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] = + mean_input_flows[1][(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] = 4.5 * Δt_allocation Ribasim.update_allocation!(model.integrator) @@ -224,9 +225,9 @@ end allocation t = 0.0 - # Set flows of sources in - mean_input_flows[(NodeID(:FlowBoundary, 58, p), NodeID(:Basin, 16, p))] = 1.0 - mean_input_flows[(NodeID(:FlowBoundary, 59, p), NodeID(:Basin, 44, p))] = 1e-3 + # Set flows of sources in subnetworks + mean_input_flows[2][(NodeID(:FlowBoundary, 58, p), NodeID(:Basin, 16, p))] = 1.0 + mean_input_flows[4][(NodeID(:FlowBoundary, 59, p), NodeID(:Basin, 44, p))] = 1e-3 # Collecting demands u = ComponentVector(; storage = zeros(length(basin.node_id))) @@ -251,24 +252,24 @@ end du = get_du(model.integrator) Ribasim.formulate_storages!(current_storage, du, u, p, t) - current_storage ≈ Float32[ - 1.0346e6, - 1.0346e6, - 1.0346e6, - 1.0346e6, - 1.0346e6, - 13.83, - 40.10, - 6.029e5, - 4641, - 2402, - 6.03995, - 928.8, - 8.017, - 10417, - 5.619, - 10417, - 4.057, + @test current_storage ≈ Float32[ + 1.0346908f6, + 1.03469f6, + 1.0346894f6, + 1.034689f6, + 1.0346888f6, + 13.833241, + 40.109993, + 187761.73, + 4641.365, + 2402.6687, + 6.039952, + 928.84283, + 8.0175905, + 10419.247, + 5.619053, + 10419.156, + 4.057502, ] end @@ -287,7 +288,7 @@ end (; user_demand, graph, allocation, basin, level_demand) = p # Initial "integrated" vertical flux - @test allocation.mean_input_flows[(NodeID(:Basin, 2, p), NodeID(:Basin, 2, p))] ≈ 1e2 + @test allocation.mean_input_flows[1][(NodeID(:Basin, 2, p), NodeID(:Basin, 2, p))] ≈ 1e2 Ribasim.solve!(model) @@ -614,5 +615,5 @@ end # Given a max_level of Inf, the basin capacity is 0.0 because it is not possible for the basin level to be > Inf @test Ribasim.get_basin_capacity(allocation_models[1], u, p, t, basin.node_id[1]) == 0.0 @test Ribasim.get_basin_capacity(allocation_models[1], u, p, t, basin.node_id[2]) == 0.0 - @test Ribasim.get_basin_capacity(allocation_models[1], u, p, t, basin.node_id[3]) == 0.0 + @test Ribasim.get_basin_capacity(allocation_models[2], u, p, t, basin.node_id[3]) == 0.0 end