Merge branch 'combine_flow_and_control_graphs' into refactor_allocati…

…on_with_metagraph

Manifest.toml

@@ -498,6 +498,12 @@ git-tree-sha1 = "b12f05108e405dadcc2aff0008db7f831374e051"
 uuid = "29a986be-02c6-4525-aec4-84b980013641"
 version = "2.0.0"
 
+[[deps.FileIO]]
+deps = ["Pkg", "Requires", "UUIDs"]
+git-tree-sha1 = "299dc33549f68299137e51e6d49a13b5b1da9673"
+uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+version = "1.16.1"
+
 [[deps.FileWatching]]
 uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
 
@@ -669,6 +675,12 @@ git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
 uuid = "82899510-4779-5014-852e-03e436cf321d"
 version = "1.0.0"
 
+[[deps.JLD2]]
+deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "Pkg", "PrecompileTools", "Printf", "Reexport", "Requires", "TranscodingStreams", "UUIDs"]
+git-tree-sha1 = "9bbb5130d3b4fa52846546bca4791ecbdfb52730"
+uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+version = "0.4.38"
+
 [[deps.JLLWrappers]]
 deps = ["Artifacts", "Preferences"]
 git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca"
@@ -899,6 +911,12 @@ deps = ["Artifacts", "Libdl"]
 uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
 version = "2.28.2+1"
 
+[[deps.MetaGraphsNext]]
+deps = ["Graphs", "JLD2", "SimpleTraits"]
+git-tree-sha1 = "8dd4f3f8a643d53e61ff9115749f522c35a38f3f"
+uuid = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+version = "0.6.0"
+
 [[deps.Missings]]
 deps = ["DataAPI"]
 git-tree-sha1 = "f66bdc5de519e8f8ae43bdc598782d35a25b1272"
@@ -1171,7 +1189,7 @@ uuid = "ae029012-a4dd-5104-9daa-d747884805df"
 version = "1.3.0"
 
 [[deps.Ribasim]]
-deps = ["Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimeZones", "TimerOutputs", "TranscodingStreams"]
+deps = ["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.4.0"

core/Project.toml

@@ -17,6 +17,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -26,6 +27,7 @@ JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 Legolas = "741b9549-f6ed-4911-9fbf-4a1c0c97f0cd"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 SQLite = "0aa819cd-b072-5ff4-a722-6bc24af294d9"
@@ -55,16 +57,18 @@ Dates = "<0.0.1,1"
 Dictionaries = "0.3.25"
 DiffEqCallbacks = "2.29.1"
 Documenter = "0.27,1"
+EnumX = "1.0"
 FiniteDiff = "2.21"
 ForwardDiff = "0.10"
-Graphs = "1.6"
+Graphs = "1.9"
 HiGHS = "1.7"
 IOCapture = "0.2"
 IterTools = "1.4"
 JuMP = "1.15"
 Legolas = "0.5"
 Logging = "<0.0.1,1"
 LoggingExtras = "1"
+MetaGraphsNext = "0.6"
 OrdinaryDiffEq = "6.7"
 PreallocationTools = "0.4"
 ReTestItems = "1.20"

core/src/Ribasim.jl

@@ -29,8 +29,9 @@ using DataInterpolations: LinearInterpolation, derivative
 using Dates
 using DBInterface: execute, prepare
 using Dictionaries: Indices, Dictionary, gettoken, dictionary
-using ForwardDiff: pickchunksize
 using DiffEqCallbacks
+using EnumX
+using ForwardDiff: pickchunksize
 using Graphs:
     add_edge!,
     adjacency_matrix,
@@ -46,6 +47,8 @@ using Graphs:
 using Legolas: Legolas, @schema, @version, validate, SchemaVersion, declared
 using Logging: current_logger, min_enabled_level, with_logger
 using LoggingExtras: EarlyFilteredLogger, LevelOverrideLogger
+using MetaGraphsNext:
+    MetaGraphsNext, MetaGraph, label_for, labels, outneighbor_labels, inneighbor_labels
 using OrdinaryDiffEq
 using PreallocationTools: DiffCache, FixedSizeDiffCache, get_tmp
 using SciMLBase

core/src/allocation.jl

@@ -1,30 +1,47 @@
+"""
+Temporary solution before 'edge_ids_flow_inv' is no longer needed.
+"""
+function get_edge_ids_flow_inv(graph::MetaGraph)
+    edge_ids_flow_inv = Dict{Int, Tuple{NodeID, NodeID}}()
+    for e in edges(graph)
+        id_src = label_for(graph, e.src)
+        id_dst = label_for(graph, e.dst)
+        edge_metadata = graph[id_src, id_dst]
+        if edge_metadata.type == EdgeType.flow
+            edge_ids_flow_inv[edge_metadata.id.value] = (id_src, id_dst)
+        end
+    end
+    return edge_ids_flow_inv
+end
+
 """
 Get:
 - The mapping from subnetwork node IDs to allocation graph node IDs
 - The mapping from allocation graph source node IDs to subnetwork source edge IDs
 """
 function get_node_id_mapping(
     p::Parameters,
-    subnetwork_node_ids::Vector{Int},
+    subnetwork_node_ids::Vector{NodeID},
     source_edge_ids::Vector{Int},
 )
-    (; lookup, connectivity) = p
-    (; graph_flow, edge_ids_flow_inv) = connectivity
+    (; connectivity) = p
+    (; graph) = connectivity
 
     # Mapping node_id => (allocgraph_node_id, type) where such a correspondence exists;
     # allocgraph_node_type in [:user, :junction, :basin, :source]
-    node_id_mapping = Dict{Int, Tuple{Int, Symbol}}()
+    node_id_mapping = Dict{NodeID, Tuple{Int, Symbol}}()
 
     # Determine the number of nodes in the allocgraph
     n_allocgraph_nodes = 0
     for subnetwork_node_id in subnetwork_node_ids
         add_allocgraph_node = false
-        node_type = lookup[subnetwork_node_id]
+        node_type = graph[subnetwork_node_id].type
 
         if node_type in [:user, :basin]
             add_allocgraph_node = true
 
-        elseif length(all_neighbors(graph_flow, subnetwork_node_id)) > 2
+        elseif length(all_neighbor_labels_type(graph, subnetwork_node_id, EdgeType.flow)) >
+               2
             # Each junction (that is, a node with more than 2 neighbors)
             # in the subnetwork gets an allocgraph node
             add_allocgraph_node = true
@@ -37,6 +54,9 @@ function get_node_id_mapping(
         end
     end
 
+    # Temporary solution!
+    edge_ids_flow_inv = get_edge_ids_flow_inv(graph)
+
     # Add nodes in the allocation graph for nodes connected in the problem to the source edges
     # One of these nodes can be outside the subnetwork, as long as the edge
     # connects to the subnetwork
@@ -71,21 +91,24 @@ the subnetwork.
 """
 function find_allocation_graph_edges!(
     graph_allocation::DiGraph{Int},
-    node_id_mapping::Dict{Int, Tuple{Int, Symbol}},
+    node_id_mapping::Dict{NodeID, Tuple{Int, Symbol}},
     p::Parameters,
-    subnetwork_node_ids::Vector{Int},
-)::Tuple{Vector{Vector{Int}}, SparseMatrixCSC{Float64, Int}}
+    subnetwork_node_ids::Vector{NodeID},
+)::Tuple{Vector{Vector{NodeID}}, SparseMatrixCSC{Float64, Int}}
     (; connectivity, user) = p
-    (; graph_flow) = connectivity
+    (; graph) = connectivity
 
-    allocgraph_edges_composite = Vector{Int}[]
+    allocgraph_edges_composite = Vector{NodeID}[]
     n_allocgraph_nodes = nv(graph_allocation)
     capacity = spzeros(n_allocgraph_nodes, n_allocgraph_nodes)
 
     for subnetwork_node_id in subnetwork_node_ids
-        subnetwork_inneighbor_ids = inneighbors(graph_flow, subnetwork_node_id)
-        subnetwork_outneighbor_ids = outneighbors(graph_flow, subnetwork_node_id)
-        subnetwork_neighbor_ids = all_neighbors(graph_flow, subnetwork_node_id)
+        subnetwork_inneighbor_ids =
+            inneighbor_labels_type(graph, subnetwork_node_id, EdgeType.flow)
+        subnetwork_outneighbor_ids =
+            outneighbor_labels_type(graph, subnetwork_node_id, EdgeType.flow)
+        subnetwork_neighbor_ids =
+            all_neighbor_labels_type(graph, subnetwork_node_id, EdgeType.flow)
 
         if subnetwork_node_id in keys(node_id_mapping)
             if subnetwork_node_id ∉ user.node_id
@@ -152,19 +175,18 @@ For the composite allocgraph edges:
 function process_allocation_graph_edges!(
     graph_allocation::DiGraph{Int},
     capacity::SparseMatrixCSC{Float64, Int},
-    allocgraph_edges_composite::Vector{Vector{Int}},
-    node_id_mapping::Dict{Int, Tuple{Int, Symbol}},
+    allocgraph_edges_composite::Vector{Vector{NodeID}},
+    node_id_mapping::Dict{NodeID, Tuple{Int, Symbol}},
     p::Parameters,
 )::SparseMatrixCSC{Float64, Int}
-    (; connectivity, lookup) = p
-    (; graph_flow) = connectivity
-    n_allocgraph_nodes = nv(graph_allocation)
+    (; connectivity) = p
+    (; graph) = connectivity
 
     for allocgraph_edge_composite in allocgraph_edges_composite
         # Find allocgraph node connected to this edge on the first end
         allocgraph_node_id_1 = nothing
         subnetwork_neighbors_side_1 =
-            all_neighbors(graph_flow, allocgraph_edge_composite[1])
+            all_neighbor_labels_type(graph, allocgraph_edge_composite[1], EdgeType.flow)
         for subnetwork_neighbor_node_id in subnetwork_neighbors_side_1
             if subnetwork_neighbor_node_id in keys(node_id_mapping)
                 allocgraph_node_id_1 = node_id_mapping[subnetwork_neighbor_node_id][1]
@@ -181,7 +203,7 @@ function process_allocation_graph_edges!(
         # Find allocgraph node connected to this edge on the second end
         allocgraph_node_id_2 = nothing
         subnetwork_neighbors_side_2 =
-            all_neighbors(graph_flow, allocgraph_edge_composite[end])
+            all_neighbor_labels_type(graph, allocgraph_edge_composite[end], EdgeType.flow)
         for subnetwork_neighbor_node_id in subnetwork_neighbors_side_2
             if subnetwork_neighbor_node_id in keys(node_id_mapping)
                 allocgraph_node_id_2 = node_id_mapping[subnetwork_neighbor_node_id][1]
@@ -207,7 +229,7 @@ function process_allocation_graph_edges!(
         # these do not constrain the composite edge capacity
         for (subnetwork_node_id_1, subnetwork_node_id_2, subnetwork_node_id_3) in
             IterTools.partition(allocgraph_edge_composite, 3, 1)
-            node_type = lookup[subnetwork_node_id_2]
+            node_type = graph[subnetwork_node_id_2].type
             node = getfield(p, node_type)
 
             # Find flow constraints
@@ -220,7 +242,7 @@ function process_allocation_graph_edges!(
             # Find flow direction constraints
             if is_flow_direction_constraining(node)
                 subnetwork_inneighbor_node_id =
-                    only(inneighbors(graph_flow, subnetwork_node_id_2))
+                    only(inneighbor_labels_type(graph, subnetwork_node_id_2, EdgeType.flow))
 
                 if subnetwork_inneighbor_node_id == subnetwork_node_id_1
                     negative_flow = false
@@ -249,7 +271,7 @@ The source nodes must only have one outneighbor.
 """
 function valid_sources(
     graph_allocation::DiGraph{Int},
-    node_id_mapping::Dict{Int, Tuple{Int, Symbol}},
+    node_id_mapping::Dict{NodeID, Tuple{Int, Symbol}},
 )::Bool
     errors = false
 
@@ -274,7 +296,7 @@ of the allocation graph node IDs of the users that do not have this problem.
 function avoid_using_own_returnflow!(
     graph_allocation::DiGraph{Int},
     allocgraph_node_ids_user::Vector{Int},
-    node_id_mapping_inverse::Dict{Int, Tuple{Int, Symbol}},
+    node_id_mapping_inverse::Dict{Int, Tuple{NodeID, Symbol}},
 )::Vector{Int}
     allocgraph_node_ids_user_with_returnflow = Int[]
     for allocgraph_node_id_user in allocgraph_node_ids_user
@@ -303,15 +325,15 @@ Build the graph used for the allocation problem.
 """
 function allocation_graph(
     p::Parameters,
-    subnetwork_node_ids::Vector{Int},
+    subnetwork_node_ids::Vector{NodeID},
     source_edge_ids::Vector{Int},
 )
     # Get the subnetwork and allocation node correspondence
     node_id_mapping, source_edge_mapping =
         get_node_id_mapping(p, subnetwork_node_ids, source_edge_ids)
 
     # Invert the node id mapping to easily translate from allocgraph nodes to subnetwork nodes
-    node_id_mapping_inverse = Dict{Int, Tuple{Int, Symbol}}()
+    node_id_mapping_inverse = Dict{Int, Tuple{NodeID, Symbol}}()
     for (subnetwork_node_id, (allocgraph_node_id, node_type)) in node_id_mapping
         node_id_mapping_inverse[allocgraph_node_id] = (subnetwork_node_id, node_type)
     end
@@ -592,7 +614,7 @@ Construct the allocation problem for the current subnetwork as a JuMP.jl model.
 """
 function allocation_problem(
     config::Config,
-    node_id_mapping::Dict{Int, Tuple{Int, Symbol}},
+    node_id_mapping::Dict{NodeID, Tuple{Int, Symbol}},
     allocgraph_node_ids_user_with_returnflow::Vector{Int},
     allocgraph_edges::Vector{Edge{Int}},
     allocgraph_edge_ids_user_demand::Dict{Int, Int},
@@ -662,7 +684,7 @@ function AllocationModel(
     config::Config,
     allocation_network_id::Int,
     p::Parameters,
-    subnetwork_node_ids::Vector{Int},
+    subnetwork_node_ids::Vector{NodeID},
     source_edge_ids::Vector{Int},
     Δt_allocation::Float64,
 )::AllocationModel
@@ -781,7 +803,7 @@ function assign_allocations!(
 )::Nothing
     (; problem, allocgraph_edge_ids_user_demand, node_id_mapping_inverse) = allocation_model
     (; connectivity, user) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; record) = user
     F = problem[:F]
     flow = get_tmp(flow, 0)
@@ -794,15 +816,18 @@ function assign_allocations!(
         # Save allocations to record
         push!(record.time, t)
         push!(record.allocation_network_id, allocation_model.allocation_network_id)
-        push!(record.user_node_id, model_node_id)
+        push!(record.user_node_id, model_node_id.value)
         push!(record.priority, user.priorities[priority_idx])
         push!(record.demand, user.demand[user_idx][priority_idx](t))
         push!(record.allocated, allocated)
         # Note: This is now the last abstraction before the allocation update,
         # should be the average abstraction since the last allocation solve
         push!(
             record.abstracted,
-            flow[only(inneighbors(graph_flow, model_node_id)), model_node_id],
+            flow[
+                only(inneighbor_labels_type(graph, model_node_id, EdgeType.flow)),
+                model_node_id,
+            ],
         )
     end
     return nothing
@@ -813,14 +838,18 @@ Set the source flows as capacities on edges in the AG.
 """
 function set_source_flows!(allocation_model::AllocationModel, p::Parameters)::Nothing
     (; problem, source_edge_mapping) = allocation_model
-    edge_ids_flow_inv = p.connectivity.edge_ids_flow_inv
+    # Temporary solution!
+    edge_ids_flow_inv = get_edge_ids_flow_inv(p.connectivity.graph)
 
     # It is assumed that the allocation procedure does not have to be differentiated.
     flow = get_tmp(p.connectivity.flow, 0)
 
     for (allocgraph_source_node_id, subnetwork_source_edge_id) in source_edge_mapping
-        edge_ids = edge_ids_flow_inv[subnetwork_source_edge_id]
-        JuMP.set_normalized_rhs(problem[:source][allocgraph_source_node_id], flow[edge_ids])
+        node_ids = edge_ids_flow_inv[subnetwork_source_edge_id]
+        JuMP.set_normalized_rhs(
+            problem[:source][allocgraph_source_node_id],
+            flow[node_ids...],
+        )
     end
     return nothing
 end