From 9d4dea57a92ef5eab6366109e7a4956a8b66a4af Mon Sep 17 00:00:00 2001 From: Bart de Koning Date: Wed, 18 Oct 2023 12:07:40 +0200 Subject: [PATCH] Review comments adressed --- core/src/Ribasim.jl | 7 +- core/src/allocation.jl | 366 ++++++++++-------- core/src/create.jl | 4 +- core/src/solve.jl | 11 +- core/test/allocation.jl | 2 +- docs/core/allocation.qmd | 69 +++- docs/core/index.qmd | 2 +- .../ribasim_testmodels/allocation.py | 8 - 8 files changed, 267 insertions(+), 202 deletions(-) diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl index 339c32414..6041bb6f3 100644 --- a/core/src/Ribasim.jl +++ b/core/src/Ribasim.jl @@ -17,6 +17,7 @@ module Ribasim import IterTools import BasicModelInterface as BMI import HiGHS +import JuMP.Model as JuMPModel import TranscodingStreams using Arrow: Arrow, Table @@ -32,17 +33,17 @@ using ForwardDiff: pickchunksize using DiffEqCallbacks using Graphs: add_edge!, - rem_edge!, adjacency_matrix, all_neighbors, DiGraph, + Edge, edges, inneighbors, nv, outneighbors, - SimpleEdge + rem_edge! -using JuMP +using JuMP: @variable, @constraint, @objective, set_normalized_rhs, optimize!, value using Legolas: Legolas, @schema, @version, validate, SchemaVersion, declared using Logging: current_logger, min_enabled_level, with_logger using LoggingExtras: EarlyFilteredLogger, LevelOverrideLogger diff --git a/core/src/allocation.jl b/core/src/allocation.jl index b7e4572e6..fb3960ea6 100644 --- a/core/src/allocation.jl +++ b/core/src/allocation.jl @@ -31,7 +31,7 @@ Construct the JuMP.jl model for allocation. Definitions ----------- - 'subnetwork' is used to refer to the original Ribasim subnetwork; -- 'AG' is used to refer to the allocation graph. +- 'allocgraph' is used to refer to the allocation graph. Inputs ------ @@ -43,79 +43,79 @@ source_edge_ids:: The IDs of the edges in the subnetwork whose flow fill be take Outputs ------- -An AllocationModel object, see solve.jl. +An AllocationModel object. """ -function get_allocation_model( +function AllocationModel( p::Parameters, subnetwork_node_ids::Vector{Int}, source_edge_ids::Vector{Int}, Δt_allocation::Float64, )::AllocationModel - (; connectivity, user, lookup, basin) = p + (; connectivity, user, lookup) = p (; graph_flow, edge_ids_flow_inv) = connectivity errors = false - # Mapping node_id => (AG_node_id, type) where such a correspondence exists; - # AG_node_type in [:user, :junction, :basin, :source] + # Mapping node_id => (allocgraph_node_id, type) where such a correspondence exists; + # allocgraph_node_type in [:user, :junction, :basin, :source] node_id_mapping = Dict{Int, Tuple{Int, Symbol}}() - # Determine the number of nodes in the AG - n_AG_nodes = 0 + # Determine the number of nodes in the allocgraph + n_allocgraph_nodes = 0 for subnetwork_node_id in subnetwork_node_ids - add_AG_node = false + add_allocgraph_node = false node_type = lookup[subnetwork_node_id] if node_type in [:user, :basin] - add_AG_node = true + add_allocgraph_node = true elseif length(all_neighbors(graph_flow, subnetwork_node_id)) > 2 # Each junction (that is, a node with more than 2 neighbors) - # in the subnetwork gets an AG node - add_AG_node = true + # in the subnetwork gets an allocgraph node + add_allocgraph_node = true node_type = :junction end - if add_AG_node - n_AG_nodes += 1 - node_id_mapping[subnetwork_node_id] = (n_AG_nodes, node_type) + if add_allocgraph_node + n_allocgraph_nodes += 1 + node_id_mapping[subnetwork_node_id] = (n_allocgraph_nodes, node_type) end end - # Add nodes in the AG for nodes connected in the model to the source edges + # Add nodes in the allocgraph for nodes connected in the model to the source edges # One of these nodes can be outside the subnetwork, as long as the edge # connects to the subnetwork - # Source edge mapping: AG source node ID => subnetwork source edge id + # Source edge mapping: allocgraph source node ID => subnetwork source edge id source_edge_mapping = Dict{Int, Int}() for source_edge_id in source_edge_ids subnetwork_node_id_1, subnetwork_node_id_2 = edge_ids_flow_inv[source_edge_id] if subnetwork_node_id_1 ∉ keys(node_id_mapping) - n_AG_nodes += 1 - node_id_mapping[subnetwork_node_id_1] = (n_AG_nodes, :source) - source_edge_mapping[n_AG_nodes] = source_edge_id + n_allocgraph_nodes += 1 + node_id_mapping[subnetwork_node_id_1] = (n_allocgraph_nodes, :source) + source_edge_mapping[n_allocgraph_nodes] = source_edge_id else node_id_mapping[subnetwork_node_id_1][2] = :source - source_edge_mapping[n_AG_nodes] = source_edge_id + source_edge_mapping[n_allocgraph_nodes] = source_edge_id end if subnetwork_node_id_2 ∉ keys(node_id_mapping) - n_AG_nodes += 1 - node_id_mapping[subnetwork_node_id_2] = (n_AG_nodes, :junction) + n_allocgraph_nodes += 1 + node_id_mapping[subnetwork_node_id_2] = (n_allocgraph_nodes, :junction) end end - # The AG and its edge capacities - graph_allocation = DiGraph(n_AG_nodes) - capacity = spzeros(n_AG_nodes, n_AG_nodes) + # The allocgraph and its edge capacities + graph_allocation = DiGraph(n_allocgraph_nodes) + capacity = spzeros(n_allocgraph_nodes, n_allocgraph_nodes) - # The ids of the subnetwork nodes that have an equivalent in the AG + # The ids of the subnetwork nodes that have an equivalent in the allocgraph subnetwork_node_ids_represented = keys(node_id_mapping) - # This loop finds AG edges in several ways: - # - Between AG nodes whose equivalent in the subnetwork are directly connected - # - Between AG nodes whose equivalent in the subnetwork are connected + # This loop finds allocgraph edges in several ways: + # - Between allocgraph nodes whose equivalent in the subnetwork are directly connected + # - Between allocgraph nodes whose equivalent in the subnetwork are connected # with one or more non-junction nodes in between - AG_edges_composite = Vector{Int}[] + allocgraph_edges_composite = Vector{Int}[] for subnetwork_node_id in subnetwork_node_ids subnetwork_inneighbor_ids = inneighbors(graph_flow, subnetwork_node_id) subnetwork_outneighbor_ids = outneighbors(graph_flow, subnetwork_node_id) @@ -124,102 +124,112 @@ function get_allocation_model( if subnetwork_node_id in subnetwork_node_ids_represented if subnetwork_node_id ∉ user.node_id # Direct connections in the subnetwork between nodes that - # have an equivaent AG graph node + # have an equivaent allocgraph graph node for subnetwork_inneighbor_id in subnetwork_inneighbor_ids if subnetwork_inneighbor_id in subnetwork_node_ids_represented - AG_node_id_1 = node_id_mapping[subnetwork_node_id][1] - AG_node_id_2 = node_id_mapping[subnetwork_inneighbor_id][1] - add_edge!(graph_allocation, AG_node_id_2, AG_node_id_1) + allocgraph_node_id_1 = node_id_mapping[subnetwork_node_id][1] + allocgraph_node_id_2 = node_id_mapping[subnetwork_inneighbor_id][1] + add_edge!( + graph_allocation, + allocgraph_node_id_2, + allocgraph_node_id_1, + ) # These direct connections cannot have capacity constraints - capacity[AG_node_id_2, AG_node_id_1] = Inf + capacity[allocgraph_node_id_2, allocgraph_node_id_1] = Inf end end for subnetwork_outneighbor_id in subnetwork_outneighbor_ids if subnetwork_outneighbor_id in subnetwork_node_ids_represented - AG_node_id_1 = node_id_mapping[subnetwork_node_id][1] - AG_node_id_2 = node_id_mapping[subnetwork_outneighbor_id][1] - add_edge!(graph_allocation, AG_node_id_1, AG_node_id_2) + allocgraph_node_id_1 = node_id_mapping[subnetwork_node_id][1] + allocgraph_node_id_2 = node_id_mapping[subnetwork_outneighbor_id][1] + add_edge!( + graph_allocation, + allocgraph_node_id_1, + allocgraph_node_id_2, + ) if subnetwork_outneighbor_id in user.node_id # Capacity depends on user demand at a given priority - capacity[AG_node_id_1, AG_node_id_2] = Inf + capacity[allocgraph_node_id_1, allocgraph_node_id_2] = Inf else # These direct connections cannot have capacity constraints - capacity[AG_node_id_1, AG_node_id_2] = Inf + capacity[allocgraph_node_id_1, allocgraph_node_id_2] = Inf end end end end else - # Try to find an existing AG composite edge to add the current subnetwork_node_id to + # Try to find an existing allocgraph composite edge to add the current subnetwork_node_id to found_edge = false - for AG_edge_composite in AG_edges_composite - if AG_edge_composite[1] in subnetwork_neighbor_ids - pushfirst!(AG_edge_composite, subnetwork_node_id) + for allocgraph_edge_composite in allocgraph_edges_composite + if allocgraph_edge_composite[1] in subnetwork_neighbor_ids + pushfirst!(allocgraph_edge_composite, subnetwork_node_id) found_edge = true break - elseif AG_edge_composite[end] in subnetwork_neighbor_ids - push!(AG_edge_composite, subnetwork_node_id) + elseif allocgraph_edge_composite[end] in subnetwork_neighbor_ids + push!(allocgraph_edge_composite, subnetwork_node_id) found_edge = true break end end - # Start a new AG composite edge if no existing edge to append to was found + # Start a new allocgraph composite edge if no existing edge to append to was found if !found_edge - push!(AG_edges_composite, [subnetwork_node_id]) + push!(allocgraph_edges_composite, [subnetwork_node_id]) end end end - # For the composite AG edges: - # - Find out whether they are connected to AG nodes on both ends + # For the composite allocgraph edges: + # - Find out whether they are connected to allocgraph nodes on both ends # - Compute their capacity # - Find out their allowed flow direction(s) - for AG_edge_composite in AG_edges_composite - # Find AG node connected to this edge on the first end - AG_node_id_1 = nothing - subnetwork_neighbors_side_1 = all_neighbors(graph_flow, AG_edge_composite[1]) + for allocgraph_edge_composite in allocgraph_edges_composite + # Find allocgraph node connected to this edge on the first end + allocgraph_node_id_1 = nothing + subnetwork_neighbors_side_1 = + all_neighbors(graph_flow, allocgraph_edge_composite[1]) for subnetwork_neighbor_node_id in subnetwork_neighbors_side_1 if subnetwork_neighbor_node_id in subnetwork_node_ids_represented - AG_node_id_1 = node_id_mapping[subnetwork_neighbor_node_id][1] - pushfirst!(AG_edge_composite, subnetwork_neighbor_node_id) + allocgraph_node_id_1 = node_id_mapping[subnetwork_neighbor_node_id][1] + pushfirst!(allocgraph_edge_composite, subnetwork_neighbor_node_id) break end end # No connection to a max flow node found on this side, so edge is discarded - if isnothing(AG_node_id_1) + if isnothing(allocgraph_node_id_1) continue end - # Find AG node connected to this edge on the second end - AG_node_id_2 = nothing - subnetwork_neighbors_side_2 = all_neighbors(graph_flow, AG_edge_composite[end]) + # Find allocgraph node connected to this edge on the second end + allocgraph_node_id_2 = nothing + subnetwork_neighbors_side_2 = + all_neighbors(graph_flow, allocgraph_edge_composite[end]) for subnetwork_neighbor_node_id in subnetwork_neighbors_side_2 if subnetwork_neighbor_node_id in subnetwork_node_ids_represented - AG_node_id_2 = node_id_mapping[subnetwork_neighbor_node_id][1] - # Make sure this AG node is distinct from the other one - if AG_node_id_2 ≠ AG_node_id_1 - push!(AG_edge_composite, subnetwork_neighbor_node_id) + allocgraph_node_id_2 = node_id_mapping[subnetwork_neighbor_node_id][1] + # Make sure this allocgraph node is distinct from the other one + if allocgraph_node_id_2 ≠ allocgraph_node_id_1 + push!(allocgraph_edge_composite, subnetwork_neighbor_node_id) break end end end - # No connection to AG node found on this side, so edge is discarded - if isnothing(AG_node_id_2) + # No connection to allocgraph node found on this side, so edge is discarded + if isnothing(allocgraph_node_id_2) continue end - # Find capacity of this composite AG edge + # Find capacity of this composite allocgraph edge positive_flow = true negative_flow = true - AG_edge_capacity = Inf - for (i, subnetwork_node_id) in enumerate(AG_edge_composite) - # The start and end subnetwork nodes of the composite AG - # edge are now nodes that have an equivalent in the AG graph, + allocgraph_edge_capacity = Inf + for (i, subnetwork_node_id) in enumerate(allocgraph_edge_composite) + # The start and end subnetwork nodes of the composite allocgraph + # edge are now nodes that have an equivalent in the allocgraph graph, # these do not constrain the composite edge capacity - if i == 1 || i == length(AG_edge_composite) + if i == 1 || i == length(allocgraph_edge_composite) continue end node_type = lookup[subnetwork_node_id] @@ -228,7 +238,8 @@ function get_allocation_model( # Find flow constraints if is_flow_constraining(node) model_node_idx = Ribasim.findsorted(node.node_id, subnetwork_node_id) - AG_edge_capacity = min(AG_edge_capacity, node.max_flow_rate[model_node_idx]) + allocgraph_edge_capacity = + min(allocgraph_edge_capacity, node.max_flow_rate[model_node_idx]) end # Find flow direction constraints @@ -236,32 +247,32 @@ function get_allocation_model( subnetwork_inneighbor_node_id = only(inneighbors(graph_flow, subnetwork_node_id)) - if subnetwork_inneighbor_node_id == AG_edge_composite[i - 1] + if subnetwork_inneighbor_node_id == allocgraph_edge_composite[i - 1] negative_flow = false - elseif subnetwork_inneighbor_node_id == AG_edge_composite[i + 1] + elseif subnetwork_inneighbor_node_id == allocgraph_edge_composite[i + 1] positive_flow = false end end end - # Add composite AG edge(s) + # Add composite allocgraph edge(s) if positive_flow - add_edge!(graph_allocation, AG_node_id_1, AG_node_id_2) - capacity[AG_node_id_1, AG_node_id_2] = AG_edge_capacity + add_edge!(graph_allocation, allocgraph_node_id_1, allocgraph_node_id_2) + capacity[allocgraph_node_id_1, allocgraph_node_id_2] = allocgraph_edge_capacity end if negative_flow - add_edge!(graph_allocation, AG_node_id_2, AG_node_id_1) - capacity[AG_node_id_2, AG_node_id_1] = AG_edge_capacity + add_edge!(graph_allocation, allocgraph_node_id_2, allocgraph_node_id_1) + capacity[allocgraph_node_id_2, allocgraph_node_id_1] = allocgraph_edge_capacity end end # The source nodes must only have one outneighbor - for (AG_node_id, AG_node_type) in values(node_id_mapping) - if AG_node_type == :source + for (allocgraph_node_id, allocgraph_node_type) in values(node_id_mapping) + if allocgraph_node_type == :source if !( - (length(inneighbors(graph_allocation, AG_node_id)) == 0) && - (length(outneighbors(graph_allocation, AG_node_id)) == 1) + (length(inneighbors(graph_allocation, allocgraph_node_id)) == 0) && + (length(outneighbors(graph_allocation, allocgraph_node_id)) == 1) ) @error "Sources nodes in the max flow graph must have no inneighbors and 1 outneighbor." errors = true @@ -269,85 +280,93 @@ function get_allocation_model( end end - # Invert the node id mapping to easily translate from AG nodes to subnetwork nodes + # Invert the node id mapping to easily translate from allocgraph nodes to subnetwork nodes node_id_mapping_inverse = Dict{Int, Tuple{Int, Symbol}}() - for (subnetwork_node_id, (AG_node_id, node_type)) in node_id_mapping - node_id_mapping_inverse[AG_node_id] = (subnetwork_node_id, node_type) + for (subnetwork_node_id, (allocgraph_node_id, node_type)) in node_id_mapping + node_id_mapping_inverse[allocgraph_node_id] = (subnetwork_node_id, node_type) end # Remove user return flow edges that are upstream of the user itself - AG_node_ids_user = [ - AG_node_id for - (AG_node_id, node_type) in values(node_id_mapping) if node_type == :user + allocgraph_node_ids_user = [ + allocgraph_node_id for + (allocgraph_node_id, node_type) in values(node_id_mapping) if node_type == :user ] - AG_node_ids_user_with_returnflow = Int[] - for AG_node_id_user in AG_node_ids_user - AG_node_id_return_flow = only(outneighbors(graph_allocation, AG_node_id_user)) - if path_exists(graph_allocation, AG_node_id_return_flow, AG_node_id_user) - rem_edge!(graph_allocation, AG_node_id_user, AG_node_id_return_flow) - # TODO: Add to logging? - @warn "The outflow of user #$(node_id_mapping_inverse[AG_node_id_user][1]) is upstream of this user itself and thus ignored." + allocgraph_node_ids_user_with_returnflow = Int[] + for allocgraph_node_id_user in allocgraph_node_ids_user + allocgraph_node_id_return_flow = + only(outneighbors(graph_allocation, allocgraph_node_id_user)) + if path_exists( + graph_allocation, + allocgraph_node_id_return_flow, + allocgraph_node_id_user, + ) + rem_edge!( + graph_allocation, + allocgraph_node_id_user, + allocgraph_node_id_return_flow, + ) + @debug "The outflow of user #$(node_id_mapping_inverse[allocgraph_node_id_user][1]) is upstream of the user itself and thus ignored in allocation solves." else - push!(AG_node_ids_uer_with_returnflow, AG_node_id_user) + push!(allocgraph_node_ids_uer_with_returnflow, allocgraph_node_id_user) end end # Used for updating user demand and source flow constraints - AG_edges = collect(edges(graph_allocation)) - AG_edge_ids_user_demand = Int[] - AG_edge_ids_source = Int[] - for (i, AG_edge) in enumerate(AG_edges) - AG_node_type_dst = node_id_mapping_inverse[AG_edge.dst][2] - AG_node_type_src = node_id_mapping_inverse[AG_edge.src][2] - if AG_node_type_dst == :user - push!(AG_edge_ids_user_demand, i) - elseif AG_node_type_src == :source - push!(AG_edge_ids_source, i) + allocgraph_edges = collect(edges(graph_allocation)) + allocgraph_edge_ids_user_demand = Int[] + allocgraph_edge_ids_source = Int[] + for (i, allocgraph_edge) in enumerate(allocgraph_edges) + allocgraph_node_type_dst = node_id_mapping_inverse[allocgraph_edge.dst][2] + allocgraph_node_type_src = node_id_mapping_inverse[allocgraph_edge.src][2] + if allocgraph_node_type_dst == :user + push!(allocgraph_edge_ids_user_demand, i) + elseif allocgraph_node_type_src == :source + push!(allocgraph_edge_ids_source, i) end end # The JuMP.jl allocation model - model = JuMP.Model(HiGHS.Optimizer) + model = JuMPModel(HiGHS.Optimizer) # The flow variables - # The variable indices are the AG edge IDs. - n_flows = length(AG_edges) + # The variable indices are the allocgraph edge IDs. + n_flows = length(allocgraph_edges) model[:F] = @variable(model, F[1:n_flows] >= 0.0) # The user allocation variables - # The variable name indices are the AG user node IDs + # The variable name indices are the allocgraph user node IDs # The variable indices are the priorities. - for AG_edge_id_user_demand in AG_edge_ids_user_demand - AG_node_id_user = AG_edges[AG_edge_id_user_demand].dst - base_name = "A_user_$AG_node_id_user" + for allocgraph_edge_id_user_demand in allocgraph_edge_ids_user_demand + allocgraph_node_id_user = allocgraph_edges[allocgraph_edge_id_user_demand].dst + base_name = "A_user_$allocgraph_node_id_user" model[Symbol(base_name)] = @variable(model, [1:length(user.priorities)], base_name = base_name) end # The basin allocation variables - # The variable indices are the AG basin node IDs - AG_node_ids_basin = sort([ - AG_node_id for - (AG_node_id, node_type) in values(node_id_mapping) if node_type == :basin + # The variable indices are the allocgraph basin node IDs + allocgraph_node_ids_basin = sort([ + allocgraph_node_id for + (allocgraph_node_id, node_type) in values(node_id_mapping) if node_type == :basin ]) - @variable(model, A_basin[i = AG_node_ids_basin] >= 0.0) + @variable(model, A_basin[i = allocgraph_node_ids_basin] >= 0.0) # The user allocation constraints - for AG_edge_id_user_demand in AG_edge_ids_user_demand - AG_node_id_user = AG_edges[AG_edge_id_user_demand].dst - base_name = "A_user_$AG_node_id_user" + for allocgraph_edge_id_user_demand in allocgraph_edge_ids_user_demand + allocgraph_node_id_user = allocgraph_edges[allocgraph_edge_id_user_demand].dst + base_name = "A_user_$allocgraph_node_id_user" A_user = model[Symbol(base_name)] # Sum of allocations to user is total flow to user @constraint( model, - sum(A_user) == F[AG_edge_id_user_demand], - base_name = "allocation_sum[$AG_node_id_user]" + sum(A_user) == F[allocgraph_edge_id_user_demand], + base_name = "allocation_sum[$allocgraph_node_id_user]" ) # Allocation flows are non-negative @constraint(model, [p = 1:length(user.priorities)], A_user[p] >= 0) # Allocation flows are bounded from above by demands - base_name = "demand_user_$AG_node_id_user" + base_name = "demand_user_$allocgraph_node_id_user" model[Symbol(base_name)] = @constraint( model, [p = 1:length(user.priorities)], @@ -358,65 +377,69 @@ function get_allocation_model( # The basin allocation constraints (actual threshold values will be set before # each allocation solve) - # The constraint indices are the AG basin node IDs + # The constraint indices are the allocgraph basin node IDs model[:basin_allocation] = @constraint( model, - [i = AG_node_ids_basin], + [i = allocgraph_node_ids_basin], A_basin[i] <= 0.0, base_name = "basin_allocation" ) # The capacity constraints - # The constraint indices are the AG edge IDs - AG_edge_ids_finite_capacity = Int[] - for (i, AG_edge) in enumerate(AG_edges) - if !isinf(capacity[AG_edge.src, AG_edge.dst]) - push!(AG_edge_ids_finite_capacity, i) + # The constraint indices are the allocgraph edge IDs + allocgraph_edge_ids_finite_capacity = Int[] + for (i, allocgraph_edge) in enumerate(allocgraph_edges) + if !isinf(capacity[allocgraph_edge.src, allocgraph_edge.dst]) + push!(allocgraph_edge_ids_finite_capacity, i) end end model[:capacity] = @constraint( model, - [i = AG_edge_ids_finite_capacity], - F[i] <= capacity[AG_edges[i].src, AG_edges[i].dst], + [i = allocgraph_edge_ids_finite_capacity], + F[i] <= capacity[allocgraph_edges[i].src, allocgraph_edges[i].dst], base_name = "capacity" ) # The source constraints (actual threshold values will be set before # each allocation solve) - # The constraint indices are the AG source node IDs + # The constraint indices are the allocgraph source node IDs model[:source] = @constraint( model, [i = keys(source_edge_mapping)], F[findfirst( - ==(SimpleEdge(i, only(outneighbors(graph_allocation, i)))), - AG_edges, + ==(Edge(i, only(outneighbors(graph_allocation, i)))), + allocgraph_edges, )] <= 0.0, base_name = "source" ) # The user return flow constraints - # The constraint indices are AG user node IDs - AG_node_inedge_ids = Dict(i => Int[] for i in 1:n_AG_nodes) - AG_node_outedge_ids = Dict(i => Int[] for i in 1:n_AG_nodes) - for (i, AG_edge) in enumerate(AG_edges) - push!(AG_node_inedge_ids[AG_edge.dst], i) - push!(AG_node_outedge_ids[AG_edge.src], i) + # The constraint indices are allocgraph user node IDs + allocgraph_node_inedge_ids = Dict(i => Int[] for i in 1:n_allocgraph_nodes) + allocgraph_node_outedge_ids = Dict(i => Int[] for i in 1:n_allocgraph_nodes) + for (i, allocgraph_edge) in enumerate(allocgraph_edges) + push!(allocgraph_node_inedge_ids[allocgraph_edge.dst], i) + push!(allocgraph_node_outedge_ids[allocgraph_edge.src], i) end model[:return_flow] = @constraint( model, - [i = AG_node_ids_user_with_returnflow], - F[only(AG_node_outedge_ids[i])] == + [i = allocgraph_node_ids_user_with_returnflow], + F[only(allocgraph_node_outedge_ids[i])] == user.return_factor[findsorted(user.node_id, node_id_mapping_inverse[i][1])] * - F[only(AG_node_inedge_ids[i])], + F[only(allocgraph_node_inedge_ids[i])], base_name = "return_flow", ) # The flow conservation constraints - # The constraint indices are AG user node IDs + # The constraint indices are allocgraph user node IDs model[:flow_conservation] = @constraint( model, - [i = AG_node_ids_basin], - sum([F[AG_edge_id] for AG_edge_id in AG_node_outedge_ids[i]]) <= sum([F[AG_edge_id] for AG_edge_id in AG_node_inedge_ids[i]]), + [i = allocgraph_node_ids_basin], + sum([ + F[allocgraph_edge_id] for allocgraph_edge_id in allocgraph_node_outedge_ids[i] + ]) <= sum([ + F[allocgraph_edge_id] for allocgraph_edge_id in allocgraph_node_inedge_ids[i] + ]), base_name = "flow_conservation", ) @@ -424,8 +447,10 @@ function get_allocation_model( # The objective function allocation_user_weights = 1 ./ (2 .^ (1:length(user.priorities))) - allocation_user_variables = - [model[Symbol("A_user_$AG_node_id_user")] for AG_node_id_user in AG_node_ids_user] + allocation_user_variables = [ + model[Symbol("A_user_$allocgraph_node_id_user")] for + allocgraph_node_id_user in allocgraph_node_ids_user + ] @objective( model, Max, @@ -447,25 +472,26 @@ function get_allocation_model( ) end +""" +Update the allocation optimization problem for the given subnetwork with the model state +and flows, solve the allocation problem and assign the results to the users. +""" function allocate!(p::Parameters, allocation_model::AllocationModel, t::Float64;)::Nothing (; node_id_mapping, source_edge_mapping, model) = allocation_model (; user, connectivity) = p (; priorities, demand) = user (; flow, edge_ids_flow_inv) = connectivity - # TODO: It is assumed here that the allocation procedure does not have to be differentiated. - # However, the allocation depends on the model state so that is not true. I have a suspicion - # that HiGHS will not work together with ForwardDiff.jl, so we might need an LP solver that - # is compatible with the AD package that we use. + # It is assumed that the allocation procedure does not have to be differentiated. flow = get_tmp(flow, 0) - for (subnetwork_node_id, (AG_node_id, AG_node_type)) in node_id_mapping - if AG_node_type == :user + for (subnetwork_node_id, (allocgraph_node_id, allocgraph_node_type)) in node_id_mapping + if allocgraph_node_type == :user # Set the user demands at the current time as the upper bound to # the allocations to the users node_idx = findsorted(user.node_id, subnetwork_node_id) demand = user.demand[node_idx] - base_name = "demand_user_$AG_node_id" + base_name = "demand_user_$allocgraph_node_id" constraints_demand = model[Symbol(base_name)] for priority_idx in eachindex(priorities) set_normalized_rhs( @@ -473,30 +499,32 @@ function allocate!(p::Parameters, allocation_model::AllocationModel, t::Float64; demand[priority_idx](t), ) end - elseif AG_node_type == :source + elseif allocgraph_node_type == :source # Set the source flows as the source flow upper bounds in # the allocation model - subnetwork_edge = source_edge_mapping[AG_node_id] + subnetwork_edge = source_edge_mapping[allocgraph_node_id] subnetwork_node_ids = edge_ids_flow_inv[subnetwork_edge] - constraint_source = model[:source][AG_node_id] + constraint_source = model[:source][allocgraph_node_id] set_normalized_rhs(constraint_source, flow[subnetwork_node_ids]) - elseif AG_node_type == :basin + elseif allocgraph_node_type == :basin # TODO: Compute basin flow from vertical fluxes and basin volume. # Set as basin demand if the net flow is negative, set as source # in the flow_conservation constraints if the net flow is positive. - elseif AG_node_type == :junction + elseif allocgraph_node_type == :junction nothing else - error("Got unsupported allocation graph node type $AG_node_type.") + error("Got unsupported allocation graph node type $allocgraph_node_type.") end end + # Solve the allocation problem optimize!(model) - for (subnetwork_node_id, (AG_node_id, AG_node_type)) in node_id_mapping - if AG_node_type == :user + # Assign the allocations to the users + for (subnetwork_node_id, (allocgraph_node_id, allocgraph_node_type)) in node_id_mapping + if allocgraph_node_type == :user user_idx = findsorted(user.node_id, subnetwork_node_id) - base_name = "A_user_$AG_node_id" + base_name = "A_user_$allocgraph_node_id" user.allocated[user_idx] .= value.(model[Symbol(base_name)]) end end diff --git a/core/src/create.jl b/core/src/create.jl index 8543d238d..6e0c0aa01 100644 --- a/core/src/create.jl +++ b/core/src/create.jl @@ -231,7 +231,7 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity end # TODO: Create allocation models from input here - allocation_model = AllocationModel[] + allocation_models = AllocationModel[] return Connectivity( graph_flow, @@ -242,7 +242,7 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity edge_ids_control, edge_connection_types_flow, edge_connection_types_control, - allocation_model, + allocation_models, ) end diff --git a/core/src/solve.jl b/core/src/solve.jl index 9ee41a3b1..71eefec37 100644 --- a/core/src/solve.jl +++ b/core/src/solve.jl @@ -6,14 +6,13 @@ const VectorInterpolation = """ Store information for a subnetwork used for allocation. -For more information see allocation.jl. node_id: All the IDs of the nodes that are in this subnetwork node_id_mapping: Mapping Dictionary; model_node_id => AG_node_id where such a correspondence exists (all AG node ids are in the values) -node_id_mapping_inverse: The inverse of node_id_mapping, Dictionary; AG_node_id => model_node_id -Source edge mapping: AG source node ID => subnetwork source edge id -graph_max_flow: The graph used for the allocation problems +node_id_mapping_inverse: The inverse of node_id_mapping, Dictionary; AG node ID => model node ID +Source edge mapping: AG source node ID => subnetwork source edge ID +graph_allocation: The graph used for the allocation problems capacity: The capacity per edge of the allocation graph, as constrained by nodes that have a max_flow_rate model: The JuMP.jl model for solving the allocation problem Δt_allocation: The time interval between consecutive allocation solves @@ -25,7 +24,7 @@ struct AllocationModel source_edge_mapping::Dict{Int, Int} graph_allocation::DiGraph{Int} capacity::SparseMatrixCSC{Float64, Int} - model::JuMP.Model + model::JuMPModel Δt_allocation::Float64 end @@ -52,7 +51,7 @@ struct Connectivity{T} edge_ids_control::Dictionary{Tuple{Int, Int}, Int} edge_connection_type_flow::Dictionary{Int, Tuple{Symbol, Symbol}} edge_connection_type_control::Dictionary{Int, Tuple{Symbol, Symbol}} - allocation_model::Vector{AllocationModel} + allocation_models::Vector{AllocationModel} function Connectivity( graph_flow, graph_control, diff --git a/core/test/allocation.jl b/core/test/allocation.jl index 74e5c07cb..b5b72f930 100644 --- a/core/test/allocation.jl +++ b/core/test/allocation.jl @@ -21,7 +21,7 @@ using JuMP: value t = 0.0 allocation_model = - Ribasim.get_allocation_model(p, subgraph_node_ids, source_edge_ids, Δt_allocation) + Ribasim.AllocationModel(p, subgraph_node_ids, source_edge_ids, Δt_allocation) Ribasim.allocate!(p, allocation_model, t) F = value.(allocation_model.model[:F]) diff --git a/docs/core/allocation.qmd b/docs/core/allocation.qmd index eb51e67c1..a8706f8d6 100644 --- a/docs/core/allocation.qmd +++ b/docs/core/allocation.qmd @@ -2,9 +2,9 @@ title: "Allocation" --- -Allocation is the process of associating an allocated abstraction flow rate to user nodes in the model based on information about sources, user demands over various priorities, constraints introduced by nodes, local water availability and graph topology. The allocation procedure implemented in Ribasim is heavily inspired by the [maximum flow problem](https://en.wikipedia.org/wiki/Maximum_flow_problem). +Allocation is the process of assigning an allocated abstraction flow rate to user nodes in the model based on information about sources, user demands over various priorities, constraints introduced by nodes, local water availability and graph topology. The allocation procedure implemented in Ribasim is heavily inspired by the [maximum flow problem](https://en.wikipedia.org/wiki/Maximum_flow_problem). -The allocation problem is solved per subnetwork of the Ribasim model. The subnetwork is used to formulate an optimization problem which is solved using the `JuMP.jl` package. +The allocation problem is solved per subnetwork of the Ribasim model. The subnetwork is used to formulate an optimization problem with the [JuMP](https://jump.dev/JuMP.jl/stable/) package, which is solved using the [HiGHS solver](https://highs.dev/). See also the example of solving the maximum flow problem with `JuMP.jl` [here](https://jump.dev/JuMP.jl/stable/tutorials/linear/network_flows/#The-max-flow-problem). # The allocation problem @@ -18,7 +18,7 @@ The allocation problem is solved per subgraph, where a subgraph is given by a su ### Source flows -Sources are indicated a set of edges in the subnetwork +Sources are indicated by a set of edges in the subnetwork $$ E_S^\text{source} \subset \left(S \times S\right) \cap E. $$ @@ -37,7 +37,7 @@ On this page we assume that the priorities are given by all integers from $1$ to ### Vertical fluxes and local storage -Apart from the source flows denoted by edges, there are other sources of water in the subnetwork, associated with the basins in the subnetwork $B_S = B \cap S$. Firstly there is the sum of the horizontal fluxes (precipitation, evaporation, infiltration and drainage) for each basin +Apart from the source flows denoted by edges, there are other sources of water in the subnetwork, associated with the basins in the subnetwork $B_S = B \cap S$. Firstly there is the sum of the vertical fluxes (precipitation, evaporation, infiltration and drainage) for each basin $$ \phi_i(t), \quad \forall i \in B_S. $$ @@ -62,7 +62,7 @@ Both fractional flow nodes and user nodes dictate proportional relationships bet A new graph is created from the subnetwork, which we call the allocation graph. To indicate the difference between subnetwork data and allocation graph data, the allocation graph data is denoted with a hat. The allocation graph consists of: -- Nodes $\hat{V_S}$, where each basin, source and user in the subnetwork get a node in the allocation graph. Also nodes that have fractional flow outneighbors get a node in the allocation graph. The implementation makes heavy use of the node id mapping $m_S : i \mapsto \hat{i}$ to translate from subnetwork node Ds to allocation graph node IDs. Unless specified otherwise, we assume this relationship between index symbols that appear both with and without a hat. +- Nodes $\hat{V_S}$, where each basin, source and user in the subnetwork get a node in the allocation graph. Also nodes that have fractional flow outneighbors get a node in the allocation graph. The implementation makes heavy use of the node id mapping $m_S : i \mapsto \hat{i}$ to translate from subnetwork node IDs to allocation graph node IDs. Unless specified otherwise, we assume this relationship between index symbols that appear both with and without a hat. - Edges $\hat{E_S}$, where the edges in the allocation graph are given by one or more edges in the subnetwork, where those edges connect nodes in the subnetwork that have an equivalent in the allocation graph. The direction of the edges in the allocation graph is given by the direction constraints in the subnetwork. For notational convenience, we use the notation @@ -76,9 +76,9 @@ for the set of in-neighbors and out-neighbors of a node in the allocation graph ### The allocation graph capacities -The capacities of the edges of the allocation graph are collected in the sparse capacity matrix $\hat{C}_S \in \overline{\mathbb{R}}_{\ge 0}^{\hat{n}\times\hat{n}}$ where $\hat{n}$ is the number of nodes in the allocation graph. The capacities can be infinity. +The capacities of the edges of the allocation graph are collected in the sparse capacity matrix $\hat{C}_S \in \overline{\mathbb{R}}_{\ge 0}^{\hat{n}\times\hat{n}}$ where $\hat{n}$ is the number of nodes in the allocation graph. The capacities can be infinite. -The The capacities are determined in 3 different ways: +The capacities are determined in 3 different ways: - If an edge does not exist, i.e. $(\hat{i},\hat{j}) \notin \hat{E}$ for certain $1 \le \hat{i},\hat{j}\le \hat{n}$, then $(\hat{C}_S)_{\hat{i},\hat{j}} = 0$; - The capacity of the edge $\hat{e} \in \hat{E_S}$ is given by the smallest `max_flow_rate` of the nodes along the equivalent edges in the subnetwork. If there are no nodes with a `max_flow_rate`, the edge capacity is infinite; @@ -165,14 +165,59 @@ $$ {#eq-fractionalflowconstraint} - Flow sign: Furthermore there are the non-negativity constraints for the flows and allocations, see [The optimization variables](allocation.qmd#the-optimization-variables). -:::{.callout-note} -Note that $\Phi_{\hat{k}}$ can become negative, in which case @eq-flowconservationconstraint becomes impossible to satisfy. What to do in this case? -::: - ## Final notes on the allocation problem ### Users using their own return flow -If not explicitly avoided, users can use their own return flow in this allocation problem formulation. Therefore, Return flow of users is only taken into account by allocation if that return flow is downstream of the user where it comes from. That is, if there is no path in the directed allocation graph from the user outflow node back to the user. +If not explicitly avoided, users can use their own return flow in this allocation problem formulation. +Therefore, return flow of users is only taken into account by allocation if that return flow is downstream of the user where it comes from. That is, if there is no path in the directed allocation graph from the user outflow node back to the user. # Solving the allocation problem + +The text below shows a representation generated by `JuMP.jl` of an optimization as described in the previous section. + +``` +Max A_basin[2] + A_basin[3] + A_basin[5] + 0.5 A_user_4[1] + 0.25 A_user_4[2] + 0.125 A_user_4[3] + 0.5 A_user_6[1] + 0.25 A_user_6[2] + 0.125 A_user_6[3] + 0.5 A_user_1[1] + 0.25 A_user_1[2] + 0.125 A_user_1[3] +Subject to + allocation_sum[1] : -F[1] + A_user_1[1] + A_user_1[2] + A_user_1[3] == 0 + allocation_sum[4] : -F[3] + A_user_4[1] + A_user_4[2] + A_user_4[3] == 0 + allocation_sum[6] : -F[5] + A_user_6[1] + A_user_6[2] + A_user_6[3] == 0 + A_user_1[1] >= 0 + A_user_1[2] >= 0 + A_user_1[3] >= 0 + A_user_4[1] >= 0 + A_user_4[2] >= 0 + A_user_4[3] >= 0 + A_user_6[1] >= 0 + A_user_6[2] >= 0 + A_user_6[3] >= 0 + demand_user_1[1] : A_user_1[1] <= 4 + demand_user_1[2] : A_user_1[2] <= 0 + demand_user_1[3] : A_user_1[3] <= 0 + demand_user_4[1] : A_user_4[1] <= 0 + demand_user_4[2] : A_user_4[2] <= 2 + demand_user_4[3] : A_user_4[3] <= 0 + demand_user_6[1] : A_user_6[1] <= 0 + demand_user_6[2] : A_user_6[2] <= 1 + demand_user_6[3] : A_user_6[3] <= 0 + basin_allocation[2] : A_basin[2] <= 0 + basin_allocation[3] : A_basin[3] <= 0 + basin_allocation[5] : A_basin[5] <= 0 + capacity[2] : F[2] <= 3 + capacity[4] : F[4] <= 4 + source[7] : F[6] <= 4.5 + flow_conservation[2] : F[1] - F[2] <= 0 + flow_conservation[3] : F[2] + F[3] - F[4] <= 0 + flow_conservation[5] : F[4] + F[5] - F[6] <= 0 + F[1] >= 0 + F[2] >= 0 + F[3] >= 0 + F[4] >= 0 + F[5] >= 0 + F[6] >= 0 + A_basin[2] >= 0 + A_basin[3] >= 0 + A_basin[5] >= 0 +``` + +A more detailed explanation of this will follow in the future. diff --git a/docs/core/index.qmd b/docs/core/index.qmd index 8f06cd9e0..407dd2bf7 100644 --- a/docs/core/index.qmd +++ b/docs/core/index.qmd @@ -5,7 +5,7 @@ title: "Julia core" With the term "core", we mean the computational engine of Ribasim. As detailed in the [usage](usage.qmd) documentation, it is generally used as a command line tool. -The theory is described on the [equations](equations.qmd) page, and more in-depth numerical considerations are described on the [numerical considerations](numerics.qmd) page. +The theory is described on the [equations](equations.qmd) page, and more in-depth numerical considerations are described on the [numerical considerations](numerics.qmd) page. As allocation is a large and self-contained part of the Ribasim core, it is described on the separate [allocation](allocation.qmd) page. The core is implemented in the [Julia programming language](https://julialang.org/), and can be found in the [Ribasim repository](https://github.com/Deltares/Ribasim) under the diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py index ec89e8bc7..2c9b6786e 100644 --- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py +++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py @@ -535,11 +535,3 @@ def looped_subnetwork_model(): ) return model - - -if __name__ == "__main__": - import matplotlib.pyplot as plt - - model = looped_subnetwork_model() - model.plot() - plt.show()