Save allocation flows for edges from physical layer

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "a2e54351da64b7dac3109810459579d647b75bb7"
+project_hash = "212d113fb58ab2e43d8fe0283d717923e72a9a88"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"
@@ -1254,7 +1254,7 @@ version = "3.5.13"
 deps = ["Accessors", "Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "EnumX", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimeZones", "TimerOutputs", "TranscodingStreams"]
 path = "core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
-version = "0.6.0"
+version = "0.7.0"
 
 [[deps.RuntimeGeneratedFunctions]]
 deps = ["ExprTools", "SHA", "Serialization"]

Project.toml

@@ -7,8 +7,10 @@ Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 Ribasim = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"

core/src/allocation.jl

@@ -45,14 +45,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
@@ -93,7 +96,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
@@ -107,7 +110,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
@@ -227,13 +230,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
@@ -290,6 +293,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
@@ -898,6 +902,9 @@ function adjust_edge_capacities!(
     end
 end
 
+"""
+Save the allocation flows per physical edge.
+"""
 function save_allocation_flows!(
     p::Parameters,
     t::Float64,
@@ -909,13 +916,19 @@ function save_allocation_flows!(
     F = problem[:F]
 
     for allocation_edge in first(F.axes)
-        push!(allocation_record.time, t)
-        push!(allocation_record.edge_id, graph[allocation_edge...].id)
-        push!(allocation_record.from_node_id, allocation_edge[1])
-        push!(allocation_record.to_node_id, allocation_edge[2])
-        push!(allocation_record.allocation_network_id, allocation_network_id)
-        push!(allocation_record.priority, priority)
-        push!(allocation_record.flow, JuMP.value(F[allocation_edge]))
+        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!(allocation_record.time, t)
+            push!(allocation_record.edge_id, edge_metadata.id)
+            push!(allocation_record.from_node_id, node_ids[i])
+            push!(allocation_record.to_node_id, node_ids[i + 1])
+            push!(allocation_record.allocation_network_id, allocation_network_id)
+            push!(allocation_record.priority, priority)
+            push!(allocation_record.flow, flow)
+        end
     end
     return nothing
 end

core/src/solve.jl

@@ -42,6 +42,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
@@ -50,6 +52,7 @@ struct EdgeMetadata
     from_id::NodeID
     to_id::NodeID
     allocation_flow::Bool
+    node_ids::Vector{NodeID}
 end
 
 abstract type AbstractParameterNode end

core/src/utils.jl

@@ -86,8 +86,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

core/test/allocation_test.jl

@@ -11,9 +11,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_models[1]
     Ribasim.allocate!(p, allocation_model, 0.0)

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,18 @@
             ),
             (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,
+            ),
+            (DateTime, Int, Int, Int, Int, Int, Float64),
+        )
         @test Tables.schema(subgrid) ==
               Tables.Schema((:time, :subgrid_id, :subgrid_level), (DateTime, Int, Float64))
     end

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