diff --git a/core/src/allocation.jl b/core/src/allocation.jl index f8248bbc3..d49cdf0da 100644 --- a/core/src/allocation.jl +++ b/core/src/allocation.jl @@ -82,14 +82,17 @@ If the edge does not exist, it is created. """ function indicate_allocation_flow!( graph::MetaGraph, - id_src::NodeID, - id_dst::NodeID, + node_ids::AbstractVector{NodeID}, )::Nothing + id_src = first(node_ids) + id_dst = last(node_ids) + if !haskey(graph, id_src, id_dst) - edge_metadata = EdgeMetadata(0, EdgeType.none, 0, id_src, id_dst, true) + edge_metadata = EdgeMetadata(0, EdgeType.none, 0, id_src, id_dst, true, node_ids) else edge_metadata = graph[id_src, id_dst] edge_metadata = @set edge_metadata.allocation_flow = true + edge_metadata = @set edge_metadata.node_ids = node_ids end graph[id_src, id_dst] = edge_metadata return nothing @@ -130,7 +133,7 @@ function find_allocation_graph_edges!( if is_allocation_source(graph, node_id, inneighbor_id) continue end - indicate_allocation_flow!(graph, inneighbor_id, node_id) + indicate_allocation_flow!(graph, [inneighbor_id, node_id]) push!(edge_ids, (inneighbor_id, node_id)) # These direct connections cannot have capacity constraints capacity[node_id, inneighbor_id] = Inf @@ -144,7 +147,7 @@ function find_allocation_graph_edges!( if is_allocation_source(graph, outneighbor_id, node_id) continue end - indicate_allocation_flow!(graph, node_id, outneighbor_id) + indicate_allocation_flow!(graph, [node_id, outneighbor_id]) push!(edge_ids, (node_id, outneighbor_id)) # if subnetwork_outneighbor_id in user.node_id: Capacity depends on user demand at a given priority # else: These direct connections cannot have capacity constraints @@ -264,13 +267,13 @@ function process_allocation_graph_edges!( # Add composite allocation graph edge(s) if positive_flow - indicate_allocation_flow!(graph, node_id_1, node_id_2) + indicate_allocation_flow!(graph, edge_composite) capacity[node_id_1, node_id_2] = edge_capacity push!(edge_ids, (node_id_1, node_id_2)) end if negative_flow - indicate_allocation_flow!(graph, node_id_2, node_id_1) + indicate_allocation_flow!(graph, reverse(edge_composite)) capacity[node_id_2, node_id_1] = edge_capacity push!(edge_ids, (node_id_2, node_id_1)) end @@ -328,6 +331,7 @@ function avoid_using_own_returnflow!(p::Parameters, allocation_network_id::Int): edge_metadata = graph[node_id_user, node_id_return_flow] graph[node_id_user, node_id_return_flow] = @set edge_metadata.allocation_flow = false + empty!(edge_metadata.node_ids) delete!(edge_ids, (node_id_user, node_id_return_flow)) @debug "The outflow of user $node_id_user is upstream of the user itself and thus ignored in allocation solves." end @@ -1067,6 +1071,40 @@ function adjust_edge_capacities!( end end +""" +Save the allocation flows per physical edge. +""" +function save_allocation_flows!( + p::Parameters, + t::Float64, + allocation_model::AllocationModel, + priority::Int, + collect_demands::Bool, +)::Nothing + (; problem, allocation_network_id) = allocation_model + (; allocation, graph) = p + (; record) = allocation + F = problem[:F] + + for allocation_edge in first(F.axes) + flow = JuMP.value(F[allocation_edge]) + edge_metadata = graph[allocation_edge...] + (; node_ids) = edge_metadata + + for i in eachindex(node_ids)[1:(end - 1)] + push!(record.time, t) + push!(record.edge_id, edge_metadata.id) + push!(record.from_node_id, node_ids[i]) + push!(record.to_node_id, node_ids[i + 1]) + push!(record.allocation_network_id, allocation_network_id) + push!(record.priority, priority) + push!(record.flow, flow) + push!(record.collect_demands, collect_demands) + end + end + return nothing +end + """ Update the allocation optimization problem for the given subnetwork with the problem state and flows, solve the allocation problem and assign the results to the users. @@ -1124,5 +1162,14 @@ function allocate!( # Assign the allocations to the users for this priority assign_allocations!(allocation_model, p, t, priority_idx; collect_demands) + + # Save the flows over all edges in the subnetwork + save_allocation_flows!( + p, + t, + allocation_model, + priorities[priority_idx], + collect_demands, + ) end end diff --git a/core/src/bmi.jl b/core/src/bmi.jl index c014dad0d..8b6df4d77 100644 --- a/core/src/bmi.jl +++ b/core/src/bmi.jl @@ -163,6 +163,11 @@ function BMI.finalize(model::Model)::Model path = results_path(config, RESULTS_FILENAME.allocation) write_arrow(path, table, compress) + # allocation flow + table = allocation_flow_table(model) + path = results_path(config, RESULTS_FILENAME.allocation_flow) + write_arrow(path, table, compress) + # exported levels table = subgrid_level_table(model) path = results_path(config, RESULTS_FILENAME.subgrid_levels) diff --git a/core/src/consts.jl b/core/src/consts.jl index 50126b04f..7c650796b 100644 --- a/core/src/consts.jl +++ b/core/src/consts.jl @@ -3,5 +3,6 @@ const RESULTS_FILENAME = ( flow = "flow.arrow", control = "control.arrow", allocation = "allocation.arrow", + allocation_flow = "allocation_flow.arrow", subgrid_levels = "subgrid_levels.arrow", ) diff --git a/core/src/create.jl b/core/src/create.jl index 7c5494117..0c571e7af 100644 --- a/core/src/create.jl +++ b/core/src/create.jl @@ -814,6 +814,16 @@ function Parameters(db::DB, config::Config)::Parameters Vector{Tuple{NodeID, NodeID}}[], Dict{Tuple{NodeID, NodeID}, Float64}(), Dict{Tuple{NodeID, NodeID}, Float64}(), + (; + time = Float64[], + edge_id = Int[], + from_node_id = Int[], + to_node_id = Int[], + allocation_network_id = Int[], + priority = Int[], + flow = Float64[], + collect_demands = BitVector(), + ), ) if !valid_edges(graph) diff --git a/core/src/io.jl b/core/src/io.jl index 19b99e986..55482e0e5 100644 --- a/core/src/io.jl +++ b/core/src/io.jl @@ -266,6 +266,35 @@ function allocation_table( ) end +function allocation_flow_table( + model::Model, +)::@NamedTuple{ + time::Vector{DateTime}, + edge_id::Vector{Int}, + from_node_id::Vector{Int}, + to_node_id::Vector{Int}, + allocation_network_id::Vector{Int}, + priority::Vector{Int}, + flow::Vector{Float64}, + collect_demands::BitVector, +} + (; config) = model + (; record) = model.integrator.p.allocation + + time = datetime_since.(record.time, config.starttime) + + return (; + time, + record.edge_id, + record.from_node_id, + record.to_node_id, + record.allocation_network_id, + record.priority, + record.flow, + record.collect_demands, + ) +end + function subgrid_level_table( model::Model, )::@NamedTuple{ diff --git a/core/src/solve.jl b/core/src/solve.jl index d766a86d5..317453640 100644 --- a/core/src/solve.jl +++ b/core/src/solve.jl @@ -28,6 +28,8 @@ allocation models: The allocation models for the main network and subnetworks co allocation_network_ids main_network_connections: (from_id, to_id) from the main network to the subnetwork per subnetwork subnetwork_demands: The demand of an edge from the main network to a subnetwork +record: A record of all flows computed by allocation optimization, eventually saved to + output file """ struct Allocation allocation_network_ids::Vector{Int} @@ -35,6 +37,16 @@ struct Allocation main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}} subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} + record::@NamedTuple{ + time::Vector{Float64}, + edge_id::Vector{Int}, + from_node_id::Vector{Int}, + to_node_id::Vector{Int}, + allocation_network_id::Vector{Int}, + priority::Vector{Int}, + flow::Vector{Float64}, + collect_demands::BitVector, + } end @enumx EdgeType flow control none @@ -58,6 +70,8 @@ allocation_network_id_source: ID of allocation network where this edge is a sour from_id: the node ID of the source node to_id: the node ID of the destination node allocation_flow: whether this edge has a flow in an allocation graph +node_ids: if this edge has allocation flow, these are all the + nodes from the physical layer this edge consists of """ struct EdgeMetadata id::Int @@ -66,6 +80,7 @@ struct EdgeMetadata from_id::NodeID to_id::NodeID allocation_flow::Bool + node_ids::Vector{NodeID} end abstract type AbstractParameterNode end diff --git a/core/src/utils.jl b/core/src/utils.jl index f2f743da4..b90f08237 100644 --- a/core/src/utils.jl +++ b/core/src/utils.jl @@ -93,8 +93,15 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra if ismissing(allocation_network_id) allocation_network_id = 0 end - edge_metadata = - EdgeMetadata(fid, edge_type, allocation_network_id, id_src, id_dst, false) + edge_metadata = EdgeMetadata( + fid, + edge_type, + allocation_network_id, + id_src, + id_dst, + false, + NodeID[], + ) graph[id_src, id_dst] = edge_metadata if edge_type == EdgeType.flow flow_counter += 1 diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl index 5e1eaf0de..35632e181 100644 --- a/core/test/allocation_test.jl +++ b/core/test/allocation_test.jl @@ -10,9 +10,19 @@ db = SQLite.DB(db_path) p = Ribasim.Parameters(db, cfg) + graph = p.graph close(db) - graph = p.graph + # Test compound allocation edge data + for edge_metadata in values(graph.edge_data) + if edge_metadata.allocation_flow + @test first(edge_metadata.node_ids) == edge_metadata.from_id + @test last(edge_metadata.node_ids) == edge_metadata.to_id + else + @test isempty(edge_metadata.node_ids) + end + end + Ribasim.set_flow!(graph, NodeID(1), NodeID(2), 4.5) # Source flow allocation_model = p.allocation.allocation_models[1] Ribasim.allocate!(p, allocation_model, 0.0) diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl index b2d992454..bc93285f8 100644 --- a/core/test/run_models_test.jl +++ b/core/test/run_models_test.jl @@ -19,12 +19,15 @@ basin_bytes = read(normpath(dirname(toml_path), "results/basin.arrow")) control_bytes = read(normpath(dirname(toml_path), "results/control.arrow")) allocation_bytes = read(normpath(dirname(toml_path), "results/allocation.arrow")) + allocation_flow_bytes = + read(normpath(dirname(toml_path), "results/allocation_flow.arrow")) subgrid_bytes = read(normpath(dirname(toml_path), "results/subgrid_levels.arrow")) flow = Arrow.Table(flow_bytes) basin = Arrow.Table(basin_bytes) control = Arrow.Table(control_bytes) allocation = Arrow.Table(allocation_bytes) + allocation_flow = Arrow.Table(allocation_flow_bytes) subgrid = Arrow.Table(subgrid_bytes) @testset "Schema" begin @@ -52,6 +55,19 @@ ), (DateTime, Int, Int, Int, Float64, Float64, Float64), ) + @test Tables.schema(allocation_flow) == Tables.Schema( + ( + :time, + :edge_id, + :from_node_id, + :to_node_id, + :allocation_network_id, + :priority, + :flow, + :collect_demands, + ), + (DateTime, Int, Int, Int, Int, Int, Float64, Bool), + ) @test Tables.schema(subgrid) == Tables.Schema((:time, :subgrid_id, :subgrid_level), (DateTime, Int, Float64)) end diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl index 62d8cfeb7..65040d70b 100644 --- a/core/test/validation_test.jl +++ b/core/test/validation_test.jl @@ -82,7 +82,7 @@ end function set_edge_metadata!(id_1, id_2, edge_type) graph[NodeID(id_1), NodeID(id_2)] = - EdgeMetadata(0, edge_type, 0, NodeID(id_1), NodeID(id_2), false) + EdgeMetadata(0, edge_type, 0, NodeID(id_1), NodeID(id_2), false, NodeID[]) return nothing end @@ -174,7 +174,7 @@ end function set_edge_metadata!(id_1, id_2, edge_type) graph[NodeID(id_1), NodeID(id_2)] = - EdgeMetadata(0, edge_type, 0, NodeID(id_1), NodeID(id_2), false) + EdgeMetadata(0, edge_type, 0, NodeID(id_1), NodeID(id_2), false, NodeID[]) return nothing end diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py index 0e556823e..a0bfa342d 100644 --- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py +++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py @@ -1731,7 +1731,7 @@ def main_network_with_subnetworks_model(): tabulated_rating_curve=rating_curve, allocation=allocation, starttime="2020-01-01 00:00:00", - endtime="2021-01-01 00:00:00", + endtime="2020-03-01 00:00:00", ) return model