Fix allocation flow output (#1355)

Fixes #1206.

---------

Co-authored-by: Martijn Visser <[email protected]>

core/src/allocation_optim.jl

@@ -837,23 +837,51 @@ function save_allocation_flows!(
     F_basin_out = problem[:F_basin_out]
 
     # Edge flows
-    for allocation_edge in first(F.axes)
+    allocation_edge_idx = 1
+    allocation_edges = first(F.axes)
+    n_allocation_edges = length(allocation_edges)
+
+    while allocation_edge_idx <= n_allocation_edges
+        allocation_edge = allocation_edges[allocation_edge_idx]
         flow_rate = JuMP.value(F[allocation_edge])
+
+        # Check whether the next allocation edge is the reverse of the current
+        # allocation edge
+        if allocation_edge_idx < n_allocation_edges &&
+           allocation_edges[allocation_edge_idx + 1] == reverse(allocation_edge)
+            # Combine the flow rates of bidirectional allocation edges
+            allocation_edge_idx += 1
+            flow_rate -= JuMP.value(F[allocation_edges[allocation_edge_idx]])
+        end
+
         edge_metadata = graph[allocation_edge...]
         (; node_ids) = edge_metadata
 
         for i in eachindex(node_ids)[1:(end - 1)]
+            # Check in which direction this edge in the physical layer exists
+            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
+
+        allocation_edge_idx += 1
     end
 
     # Basin flows

core/test/allocation_test.jl

@@ -196,6 +196,7 @@ end
     using Ribasim: NodeID, OptimizationType
     using ComponentArrays: ComponentVector
     using JuMP
+    using DataFrames: DataFrame, ByRow, transform!
 
     toml_path = normpath(
         @__DIR__,
@@ -209,7 +210,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
@@ -249,6 +251,16 @@ end
           [0.00399999999, 0.0, 0.0]
     @test subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)] ≈ [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)
+
     @test user_demand.allocated[2, :] ≈ [4.0, 0.0, 0.0]
     @test user_demand.allocated[7, :] ≈ [0.001, 0.0, 0.0]
 end