Test for existence of edges in allocation flow record

Deltares · Apr 4, 2024 · dfa1c7e · dfa1c7e
1 parent 7bb354f
commit dfa1c7e
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 6 deletions.
diff --git 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
@@ -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