Skip to content

Commit

Permalink
Save allocation flows (#1012)
Browse files Browse the repository at this point in the history
Fixes #967.

Blocking for merge: #1006.

Note to @visr and others working on flow visualization in QGIS: although
this yields a flow output file, it is different from the flow output
file for several reasons:
- There are demand gathering flows and allocation flows (which yields 2
flow values for each edge in a subnetwork for each allocation timestep)
- (For now) an edge connecting the main network to the subnetwork yields
flows from both the main network and the subnetwork.
  • Loading branch information
SouthEndMusic authored Feb 5, 2024
1 parent f37e68d commit 75b5ff3
Show file tree
Hide file tree
Showing 11 changed files with 153 additions and 13 deletions.
61 changes: 54 additions & 7 deletions core/src/allocation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
5 changes: 5 additions & 0 deletions core/src/bmi.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions core/src/consts.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)
10 changes: 10 additions & 0 deletions core/src/create.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
29 changes: 29 additions & 0 deletions core/src/io.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
15 changes: 15 additions & 0 deletions core/src/solve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,25 @@ 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}
allocation_models::Vector{AllocationModel}
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
Expand All @@ -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
Expand All @@ -66,6 +80,7 @@ struct EdgeMetadata
from_id::NodeID
to_id::NodeID
allocation_flow::Bool
node_ids::Vector{NodeID}
end

abstract type AbstractParameterNode end
Expand Down
11 changes: 9 additions & 2 deletions core/src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 11 additions & 1 deletion core/test/allocation_test.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
16 changes: 16 additions & 0 deletions core/test/run_models_test.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions core/test/validation_test.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion python/ribasim_testmodels/ribasim_testmodels/allocation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 75b5ff3

Please sign in to comment.