diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl index 5965298ea..62543c3db 100644 --- a/core/src/allocation_optim.jl +++ b/core/src/allocation_optim.jl @@ -846,15 +846,25 @@ function save_allocation_flows!( (; node_ids) = edge_metadata for i in eachindex(node_ids)[1:(end - 1)] + if haskey(graph, node_ids[i], node_ids[i + 1]) + id_from = node_ids[i] + id_to = node_ids[i + 1] + flow_rate_signed = flow_rate + else + id_from = node_ids[i + 1] + id_to = node_ids[i] + flow_rate_signed = -flow_rate + end + push!(record_flow.time, t) push!(record_flow.edge_id, edge_metadata.id) - push!(record_flow.from_node_type, string(node_ids[i].type)) - push!(record_flow.from_node_id, Int32(node_ids[i])) - push!(record_flow.to_node_type, string(node_ids[i + 1].type)) - push!(record_flow.to_node_id, Int32(node_ids[i + 1])) + push!(record_flow.from_node_type, string(id_from.type)) + push!(record_flow.from_node_id, Int32(id_from)) + push!(record_flow.to_node_type, string(id_to.type)) + push!(record_flow.to_node_id, Int32(id_to)) push!(record_flow.subnetwork_id, allocation_network_id) push!(record_flow.priority, priority) - push!(record_flow.flow_rate, flow_rate) + push!(record_flow.flow_rate, flow_rate_signed) push!(record_flow.optimization_type, string(optimization_type)) end end diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl index 80253447b..c521bbabd 100644 --- a/core/test/allocation_test.jl +++ b/core/test/allocation_test.jl @@ -202,6 +202,7 @@ end using Ribasim: NodeID, OptimizationType using ComponentArrays: ComponentVector using JuMP + using DataFrames: DataFrame, ByRow, transform! toml_path = normpath( @__DIR__, @@ -215,7 +216,8 @@ end close(db) (; allocation, user_demand, graph, basin) = p - (; allocation_models, subnetwork_demands, subnetwork_allocateds) = allocation + (; allocation_models, subnetwork_demands, subnetwork_allocateds, record_flow) = + allocation t = 0.0 # Collecting demands @@ -257,6 +259,16 @@ end @test user_demand.allocated[2] ≈ [4.0, 0.0, 0.0] @test user_demand.allocated[7] ≈ [0.001, 0.0, 0.0] + + # Test for existence of edges in allocation flow record + allocation_flow = DataFrame(record_flow) + transform!( + allocation_flow, + [:from_node_type, :from_node_id, :to_node_type, :to_node_id] => + ByRow((a, b, c, d) -> haskey(graph, NodeID(a, b), NodeID(c, d))) => + :edge_exists, + ) + @test all(allocation_flow.edge_exists) end @testitem "subnetworks with sources" begin