From 74e343bde63227cd3c4bb6ff7d1cdb09c01f63f4 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 15 Nov 2023 14:46:22 +0100
Subject: [PATCH 1/8] Towards using a metagraph

---
 Manifest.toml          |  22 ++++++-
 core/Project.toml      |   3 +-
 core/src/Ribasim.jl    |   5 +-
 core/src/bmi.jl        |   5 +-
 core/src/create.jl     |  87 +++++++++++++-------------
 core/src/solve.jl      | 138 +++++++++++++++++------------------------
 core/src/utils.jl      |  93 ++++++++++++++++++---------
 core/src/validation.jl |  89 +++++++++++++++-----------
 8 files changed, 244 insertions(+), 198 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index 054d6da8c..44f9ef131 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0-rc1"
 manifest_format = "2.0"
-project_hash = "ebee755404797b896be47dd2693eb495c11ce0d0"
+project_hash = "13373d26928355ecab6ec9870e664ac00457dfd7"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "5d2e21d7b0d8c22f67483ef95ebdc39c0e6b6003"
@@ -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"
diff --git a/core/Project.toml b/core/Project.toml
index 0e950b750..630c40fda 100644
--- a/core/Project.toml
+++ b/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"
@@ -54,7 +56,6 @@ Dictionaries = "0.3.25"
 DiffEqCallbacks = "2.29.1"
 FiniteDiff = "2.21"
 ForwardDiff = "0.10"
-Graphs = "1.6"
 HiGHS = "1.7"
 IterTools = "1.4"
 JuMP = "1.15"
diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl
index 7f68deec4..b8027f058 100644
--- a/core/src/Ribasim.jl
+++ b/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
diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index 12c6abcda..26c86bbe4 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -44,8 +44,7 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
         if !valid_pid_connectivity(
             pid_control.node_id,
             pid_control.listen_node_id,
-            connectivity.graph_flow,
-            connectivity.graph_control,
+            connectivity.graph,
             basin.node_id,
             pump.node_id,
         )
@@ -53,7 +52,7 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
         end
 
         if !valid_fractional_flow(
-            connectivity.graph_flow,
+            connectivity.graph,
             fractional_flow.node_id,
             fractional_flow.control_mapping,
         )
diff --git a/core/src/create.jl b/core/src/create.jl
index e5a50ed88..f1eba3726 100644
--- a/core/src/create.jl
+++ b/core/src/create.jl
@@ -62,7 +62,7 @@ function parse_static_and_time(
 
     # The control mapping is a dictionary with keys (node_id, control_state) to a named tuple of
     # parameter values to be assigned to the node with this node_id in the case of this control_state
-    control_mapping = Dict{Tuple{Int, String}, NamedTuple}()
+    control_mapping = Dict{Tuple{NodeID, String}, NamedTuple}()
 
     push!(keys_out, :control_mapping)
     push!(vals_out, control_mapping)
@@ -124,7 +124,7 @@ function parse_static_and_time(
                 end
                 # Add the parameter values to the control mapping
                 control_state_key = coalesce(control_state, "")
-                control_mapping[(node_id, control_state_key)] =
+                control_mapping[(NodeID(node_id), control_state_key)] =
                     NamedTuple{Tuple(parameter_names)}(Tuple(parameter_values))
             end
         elseif node_id in time_node_ids
@@ -201,13 +201,19 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity
         error("Invalid edge types found.")
     end
 
-    graph_flow, edge_ids_flow, edge_connection_types_flow = create_graph(db, "flow")
-    graph_control, edge_ids_control, edge_connection_types_control =
-        create_graph(db, "control")
+    graph = create_graph(db)
 
-    edge_ids_flow_inv = Dictionary(values(edge_ids_flow), keys(edge_ids_flow))
+    # Build the sparsity structure of the flow matrix
+    flow = spzeros(nv(graph), nv(graph))
+    for e in edges(graph)
+        label_src = label_for(graph, e.src)
+        label_dst = label_for(graph, e.dst)
+        edge_metadata = graph[label_src, label_dst]
+        if edge_metadata.type == EdgeType.flow
+            flow[label_src.id, label_dst.id] = 1.0
+        end
+    end
 
-    flow = adjacency_matrix(graph_flow, Float64)
     # Add a self-loop, i.e. an entry on the diagonal, for all non-conservative node types.
     # This is used to store the gain (positive) or loss (negative) for the water balance.
     # Note that this only affects the sparsity structure.
@@ -227,17 +233,7 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity
 
     allocation_models = AllocationModel[]
 
-    return Connectivity(
-        graph_flow,
-        graph_control,
-        flow,
-        edge_ids_flow,
-        edge_ids_flow_inv,
-        edge_ids_control,
-        edge_connection_types_flow,
-        edge_connection_types_control,
-        allocation_models,
-    )
+    return Connectivity(graph, flow, allocation_models)
 end
 
 function generate_allocation_models!(p::Parameters, db::DB, config::Config)::Nothing
@@ -299,7 +295,7 @@ function LinearResistance(db::DB, config::Config)::LinearResistance
     end
 
     return LinearResistance(
-        parsed_parameters.node_id,
+        NodeID.(parsed_parameters.node_id),
         BitVector(parsed_parameters.active),
         parsed_parameters.resistance,
         parsed_parameters.control_mapping,
@@ -320,7 +316,7 @@ function TabulatedRatingCurve(db::DB, config::Config)::TabulatedRatingCurve
     end
 
     interpolations = ScalarInterpolation[]
-    control_mapping = Dict{Tuple{Int, String}, NamedTuple}()
+    control_mapping = Dict{Tuple{NodeID, String}, NamedTuple}()
     active = BitVector()
     errors = false
 
@@ -340,7 +336,7 @@ function TabulatedRatingCurve(db::DB, config::Config)::TabulatedRatingCurve
                 is_active = coalesce(first(group).active, true)
                 interpolation, is_valid = qh_interpolation(node_id, StructVector(group))
                 if !ismissing(control_state)
-                    control_mapping[(node_id, control_state)] =
+                    control_mapping[(NodeID(node_id), control_state)] =
                         (; tables = interpolation, active = is_active)
                 end
             end
@@ -368,7 +364,13 @@ function TabulatedRatingCurve(db::DB, config::Config)::TabulatedRatingCurve
         error("Errors occurred when parsing TabulatedRatingCurve data.")
     end
 
-    return TabulatedRatingCurve(node_ids, active, interpolations, time, control_mapping)
+    return TabulatedRatingCurve(
+        NodeID.(node_ids),
+        active,
+        interpolations,
+        time,
+        control_mapping,
+    )
 end
 
 function ManningResistance(db::DB, config::Config)::ManningResistance
@@ -381,7 +383,7 @@ function ManningResistance(db::DB, config::Config)::ManningResistance
     end
 
     return ManningResistance(
-        parsed_parameters.node_id,
+        NodeID.(parsed_parameters.node_id),
         BitVector(parsed_parameters.active),
         parsed_parameters.length,
         parsed_parameters.manning_n,
@@ -400,7 +402,7 @@ function FractionalFlow(db::DB, config::Config)::FractionalFlow
     end
 
     return FractionalFlow(
-        parsed_parameters.node_id,
+        NodeID.(parsed_parameters.node_id),
         parsed_parameters.fraction,
         parsed_parameters.control_mapping,
     )
@@ -431,7 +433,11 @@ function LevelBoundary(db::DB, config::Config)::LevelBoundary
         error("Errors occurred when parsing LevelBoundary data.")
     end
 
-    return LevelBoundary(node_ids, parsed_parameters.active, parsed_parameters.level)
+    return LevelBoundary(
+        NodeID.(node_ids),
+        parsed_parameters.active,
+        parsed_parameters.level,
+    )
 end
 
 function FlowBoundary(db::DB, config::Config)::FlowBoundary
@@ -468,14 +474,18 @@ function FlowBoundary(db::DB, config::Config)::FlowBoundary
         error("Errors occurred when parsing FlowBoundary data.")
     end
 
-    return FlowBoundary(node_ids, parsed_parameters.active, parsed_parameters.flow_rate)
+    return FlowBoundary(
+        NodeID.(node_ids),
+        parsed_parameters.active,
+        parsed_parameters.flow_rate,
+    )
 end
 
 function Pump(db::DB, config::Config, chunk_size::Int)::Pump
     static = load_structvector(db, config, PumpStaticV1)
     defaults = (; min_flow_rate = 0.0, max_flow_rate = Inf, active = true)
     parsed_parameters, valid = parse_static_and_time(db, config, "Pump"; static, defaults)
-    is_pid_controlled = falses(length(parsed_parameters.node_id))
+    is_pid_controlled = falses(length(NodeID.(parsed_parameters.node_id)))
 
     if !valid
         error("Errors occurred when parsing Pump data.")
@@ -489,7 +499,7 @@ function Pump(db::DB, config::Config, chunk_size::Int)::Pump
     end
 
     return Pump(
-        parsed_parameters.node_id,
+        NodeID.(parsed_parameters.node_id),
         BitVector(parsed_parameters.active),
         flow_rate,
         parsed_parameters.min_flow_rate,
@@ -504,7 +514,7 @@ function Outlet(db::DB, config::Config, chunk_size::Int)::Outlet
     defaults =
         (; min_flow_rate = 0.0, max_flow_rate = Inf, min_crest_level = -Inf, active = true)
     parsed_parameters, valid = parse_static_and_time(db, config, "Outlet"; static, defaults)
-    is_pid_controlled = falses(length(parsed_parameters.node_id))
+    is_pid_controlled = falses(length(NodeID.(parsed_parameters.node_id)))
 
     if !valid
         error("Errors occurred when parsing Outlet data.")
@@ -518,7 +528,7 @@ function Outlet(db::DB, config::Config, chunk_size::Int)::Outlet
     end
 
     return Outlet(
-        parsed_parameters.node_id,
+        NodeID.(parsed_parameters.node_id),
         BitVector(parsed_parameters.active),
         flow_rate,
         parsed_parameters.min_flow_rate,
@@ -531,7 +541,7 @@ end
 
 function Terminal(db::DB, config::Config)::Terminal
     static = load_structvector(db, config, TerminalStaticV1)
-    return Terminal(static.node_id)
+    return Terminal(NodeID.(static.node_id))
 end
 
 function Basin(db::DB, config::Config, chunk_size::Int)::Basin
@@ -562,7 +572,7 @@ function Basin(db::DB, config::Config, chunk_size::Int)::Basin
     check_no_nans(table, "Basin")
 
     return Basin(
-        Indices(node_id),
+        Indices(NodeID.(node_id)),
         precipitation,
         potential_evaporation,
         drainage,
@@ -673,9 +683,9 @@ function PidControl(db::DB, config::Config, chunk_size::Int)::PidControl
     end
 
     return PidControl(
-        node_ids,
+        NodeID.(node_ids),
         BitVector(parsed_parameters.active),
-        parsed_parameters.listen_node_id,
+        NodeID.(parsed_parameters.listen_node_id),
         parsed_parameters.target,
         pid_parameters,
         pid_error,
@@ -847,15 +857,8 @@ function Parameters(db::DB, config::Config)::Parameters
         discrete_control,
         pid_control,
         user,
-        Dict{Int, Symbol}(),
     )
-    for (fieldname, fieldtype) in zip(fieldnames(Parameters), fieldtypes(Parameters))
-        if fieldtype <: AbstractParameterNode
-            for node_id in getfield(p, fieldname).node_id
-                p.lookup[node_id] = fieldname
-            end
-        end
-    end
+
     # Allocation data structures
     if config.allocation.use_allocation
         generate_allocation_models!(p, db, config)
diff --git a/core/src/solve.jl b/core/src/solve.jl
index effa8b62a..7fd90bbc9 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -32,6 +32,19 @@ struct AllocationModel
     Δt_allocation::Float64
 end
 
+@enumx EdgeType flow control
+
+struct NodeMetadata
+    type::Symbol
+    data_index::Int
+    allocation_network_id::Int
+end
+
+struct EdgeMetadata
+    id::EdgeID
+    type::EdgeType.T
+end
+
 """
 Store the connectivity information
 
@@ -47,52 +60,23 @@ else
 end
 """
 struct Connectivity{T}
-    graph_flow::DiGraph{Int}
-    graph_control::DiGraph{Int}
+    graph::MetaGraph{
+        Int64,
+        DiGraph{Int64},
+        Ribasim.NodeID,
+        Ribasim.NodeMetadata,
+        Ribasim.EdgeMetadata,
+        Nothing,
+        MetaGraphsNext.var"#11#13",
+        Float64,
+    }
     flow::T
-    edge_ids_flow::Dictionary{Tuple{Int, Int}, Int}
-    edge_ids_flow_inv::Dictionary{Int, Tuple{Int, Int}}
-    edge_ids_control::Dictionary{Tuple{Int, Int}, Int}
-    edge_connection_type_flow::Dictionary{Int, Tuple{Symbol, Symbol}}
-    edge_connection_type_control::Dictionary{Int, Tuple{Symbol, Symbol}}
     allocation_models::Vector{AllocationModel}
-    function Connectivity(
-        graph_flow,
-        graph_control,
-        flow::T,
-        edge_ids_flow,
-        edge_ids_flow_inv,
-        edge_ids_control,
-        edge_connection_types_flow,
-        edge_connection_types_control,
-        subnetwork,
-    ) where {T}
-        invalid_networks = Vector{String}()
-
-        if !valid_edges(edge_ids_flow, edge_connection_types_flow)
-            push!(invalid_networks, "flow")
-        end
-
-        if !valid_edges(edge_ids_control, edge_connection_types_control)
-            push!(invalid_networks, "control")
-        end
-
-        if isempty(invalid_networks)
-            new{T}(
-                graph_flow,
-                graph_control,
-                flow,
-                edge_ids_flow,
-                edge_ids_flow_inv,
-                edge_ids_control,
-                edge_connection_types_flow,
-                edge_connection_types_control,
-                subnetwork,
-            )
-        else
-            invalid_networks = join(invalid_networks, ", ")
-            error("Invalid network(s): $invalid_networks")
+    function Connectivity(graph, flow::T, allocation_models) where {T}
+        if !valid_edges(graph)
+            error("Invalid connectivity.")
         end
+        return new{T}(graph, flow, allocation_models)
     end
 end
 
@@ -116,7 +100,7 @@ else
 end
 """
 struct Basin{T, C} <: AbstractParameterNode
-    node_id::Indices{Int}
+    node_id::Indices{NodeID}
     precipitation::Vector{Float64}
     potential_evaporation::Vector{Float64}
     drainage::Vector{Float64}
@@ -184,11 +168,11 @@ time: The time table used for updating the tables
 control_mapping: dictionary from (node_id, control_state) to Q(h) and/or active state
 """
 struct TabulatedRatingCurve{C} <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     tables::Vector{ScalarInterpolation}
     time::StructVector{TabulatedRatingCurveTimeV1, C, Int}
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
 end
 
 """
@@ -203,10 +187,10 @@ resistance: the resistance to flow; Q = Δh/resistance
 control_mapping: dictionary from (node_id, control_state) to resistance and/or active state
 """
 struct LinearResistance <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     resistance::Vector{Float64}
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
 end
 
 """
@@ -243,13 +227,13 @@ Requirements:
 * (profile_width == 0) xor (profile_slope == 0)
 """
 struct ManningResistance <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     length::Vector{Float64}
     manning_n::Vector{Float64}
     profile_width::Vector{Float64}
     profile_slope::Vector{Float64}
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
 end
 
 """
@@ -264,9 +248,9 @@ fraction: The fraction in [0,1] of flow the node lets through
 control_mapping: dictionary from (node_id, control_state) to fraction
 """
 struct FractionalFlow <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     fraction::Vector{Float64}
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
 end
 
 """
@@ -275,7 +259,7 @@ active: whether this node is active
 level: the fixed level of this 'infinitely big basin'
 """
 struct LevelBoundary <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     level::Vector{ScalarInterpolation}
 end
@@ -286,7 +270,7 @@ active: whether this node is active and thus contributes flow
 flow_rate: target flow rate
 """
 struct FlowBoundary <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     flow_rate::Vector{ScalarInterpolation}
 end
@@ -301,12 +285,12 @@ control_mapping: dictionary from (node_id, control_state) to target flow rate
 is_pid_controlled: whether the flow rate of this pump is governed by PID control
 """
 struct Pump{T} <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     flow_rate::T
     min_flow_rate::Vector{Float64}
     max_flow_rate::Vector{Float64}
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
     is_pid_controlled::BitVector
 
     function Pump(
@@ -344,13 +328,13 @@ control_mapping: dictionary from (node_id, control_state) to target flow rate
 is_pid_controlled: whether the flow rate of this outlet is governed by PID control
 """
 struct Outlet{T} <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     flow_rate::T
     min_flow_rate::Vector{Float64}
     max_flow_rate::Vector{Float64}
     min_crest_level::Vector{Float64}
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
     is_pid_controlled::BitVector
 
     function Outlet(
@@ -384,7 +368,7 @@ end
 node_id: node ID of the Terminal node
 """
 struct Terminal <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
 end
 
 """
@@ -399,7 +383,7 @@ logic_mapping: Dictionary: (control node ID, truth state) => control state
 record: Namedtuple with discrete control information for results
 """
 struct DiscreteControl <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     listen_feature_id::Vector{Int}
     variable::Vector{String}
     look_ahead::Vector{Float64}
@@ -427,13 +411,13 @@ pid_params: a vector interpolation for parameters changing over time.
 error: the current error; basin_target - current_level
 """
 struct PidControl{T} <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
-    listen_node_id::Vector{Int}
+    listen_node_id::Vector{NodeID}
     target::Vector{ScalarInterpolation}
     pid_params::Vector{VectorInterpolation}
     error::T
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
 end
 
 """
@@ -447,7 +431,7 @@ priorities: All used priority values. Each user has a demand for all these prior
 record: Collected data of allocation optimizations for output file.
 """
 struct User <: AbstractParameterNode
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     active::BitVector
     demand::Vector{Vector{ScalarInterpolation}}
     allocated::Vector{Vector{Float64}}
@@ -482,7 +466,6 @@ struct Parameters{T, TSparse, C1, C2}
     discrete_control::DiscreteControl
     pid_control::PidControl{T}
     user::User
-    lookup::Dict{Int, Symbol}
 end
 
 """
@@ -491,23 +474,18 @@ number of flow/control inneighbors and outneighbors
 """
 function valid_n_neighbors(p::Parameters)::Bool
     (; connectivity) = p
-    (; graph_flow, graph_control) = connectivity
+    (; graph) = connectivity
 
     errors = false
 
     for nodefield in nodefields(p)
-        errors |= !valid_n_neighbors(getfield(p, nodefield), graph_flow, graph_control)
+        errors |= !valid_n_neighbors(getfield(p, nodefield), graph)
     end
 
     return !errors
 end
 
-function valid_n_neighbors(
-    node::AbstractParameterNode,
-    graph_flow::DiGraph{Int},
-    graph_control::DiGraph{Int},
-)::Bool
-    node_id = node.node_id
+function valid_n_neighbors(node::AbstractParameterNode, graph::MetaGraph)::Bool
     node_type = typeof(node)
     node_name = nameof(node_type)
 
@@ -516,14 +494,11 @@ function valid_n_neighbors(
 
     errors = false
 
-    for id in node_id
-        for (graph, bounds, edge_type) in zip(
-            (graph_flow, graph_control),
-            (bounds_flow, bounds_control),
-            (:flow, :control),
-        )
-            n_inneighbors = length(inneighbors(graph, id))
-            n_outneighbors = length(outneighbors(graph, id))
+    for id in node.node_id
+        for (bounds, edge_type) in
+            zip((bounds_flow, bounds_control), (EdgeType.flow, EdgeType.control))
+            n_inneighbors = length(inneighbor_labels_type(graph, id, edge_type))
+            n_outneighbors = length(outneighbor_labels_type(graph, id, edge_type))
 
             if n_inneighbors < bounds.in_min
                 @error "Nodes of type $node_type must have at least $(bounds.in_min) $edge_type inneighbor(s) (got $n_inneighbors for node #$id)."
@@ -546,7 +521,6 @@ function valid_n_neighbors(
             end
         end
     end
-
     return !errors
 end
 
diff --git a/core/src/utils.jl b/core/src/utils.jl
index a440dc58d..5c97ce18f 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -12,25 +12,55 @@ function valid_edge_types(db::DB)::Bool
     return !errors
 end
 
-"Return a directed graph, and a mapping from source and target nodes to edge fid."
-function create_graph(
-    db::DB,
-    edge_type_::String,
-)::Tuple{DiGraph, Dictionary{Tuple{Int, Int}, Int}, Dictionary{Int, Tuple{Symbol, Symbol}}}
+function create_graph(db::DB)::MetaGraph
     node_rows = execute(db, "select fid, type from Node")
-    nodes = dictionary((fid => Symbol(type) for (; fid, type) in node_rows))
-    graph = DiGraph(length(nodes))
     edge_rows = execute(db, "select fid, from_node_id, to_node_id, edge_type from Edge")
-    edge_ids = Dictionary{Tuple{Int, Int}, Int}()
-    edge_connection_types = Dictionary{Int, Tuple{Symbol, Symbol}}()
-    for (; fid, from_node_id, to_node_id, edge_type) in edge_rows
-        if edge_type == edge_type_
-            add_edge!(graph, from_node_id, to_node_id)
-            insert!(edge_ids, (from_node_id, to_node_id), fid)
-            insert!(edge_connection_types, fid, (nodes[from_node_id], nodes[to_node_id]))
+    graph = MetaGraph(
+        DiGraph();
+        label_type = NodeID,
+        vertex_data_type = NodeMetadata,
+        edge_data_type = EdgeMetadata,
+        graph_data = nothing, # In theory all remaining connectivity data could be put here
+    )
+    for (i, row) in enumerate(node_rows)
+        allocation_network_id =
+            hasproperty(row, :allocation_network_id) ? row.allocation_network_id : 0
+        graph[NodeID(row.fid)] = NodeMetadata(Symbol(row.type), i, allocation_network_id)
+    end
+    for (; from_node_id, to_node_id, edge_type, fid) in edge_rows
+        if edge_type == "flow"
+            edge_type = EdgeType.flow
+        elseif edge_type == "control"
+            edge_type = EdgeType.control
+        else
+            error("Invalid edge type $edge_type.")
         end
+        graph[NodeID(from_node_id), NodeID(to_node_id)] =
+            EdgeMetadata(EdgeID(fid), edge_type)
     end
-    return graph, edge_ids, edge_connection_types
+    return graph
+end
+
+function inneighbor_labels_type(
+    graph::MetaGraph,
+    label::NodeID,
+    edge_type::EdgeType.T,
+)::Vector{NodeID}
+    return [
+        label_in for label_in in inneighbor_labels(graph, label) if
+        graph[label_in, label].type == edge_type
+    ]
+end
+
+function outneighbor_labels_type(
+    graph::MetaGraph,
+    label::NodeID,
+    edge_type::EdgeType.T,
+)::Vector{NodeID}
+    return [
+        label_out for label_out in outneighbor_labels(graph, label) if
+        graph[label, label_out].type == edge_type
+    ]
 end
 
 "Calculate a profile storage by integrating the areas over the levels"
@@ -429,7 +459,7 @@ function get_level(
 end
 
 "Get the index of an ID in a set of indices."
-function id_index(ids::Indices{Int}, id::Int)::Tuple{Bool, Int}
+function id_index(ids::Indices{NodeID}, id::NodeID)::Tuple{Bool, Int}
     # We avoid creating Dictionary here since it converts the values to a Vector,
     # leading to allocations when used with PreallocationTools's ReinterpretArrays.
     hasindex, (_, i) = gettoken(ids, id)
@@ -488,8 +518,8 @@ Check:
 - Whether look_ahead is only supplied for condition variables given by a time-series.
 """
 function valid_discrete_control(p::Parameters, config::Config)::Bool
-    (; discrete_control, connectivity, lookup) = p
-    (; graph_control) = connectivity
+    (; discrete_control, connectivity) = p
+    (; graph) = connectivity
     (; node_id, logic_mapping, look_ahead, variable, listen_feature_id) = discrete_control
 
     t_end = seconds_since(config.endtime, config.starttime)
@@ -677,11 +707,11 @@ function update_jac_prototype!(
     node::Union{LinearResistance, ManningResistance},
 )::Nothing
     (; basin, connectivity) = p
-    (; graph_flow) = connectivity
+    (; graph) = connectivity
 
     for id in node.node_id
-        id_in = only(inneighbors(graph_flow, id))
-        id_out = only(outneighbors(graph_flow, id))
+        id_in = only(inneighbor_labels_type(graph, id, EdgeType.flow))
+        id_out = only(outneighbor_labels_type(graph, id, EdgeType.flow))
 
         has_index_in, idx_in = id_index(basin.node_id, id_in)
         has_index_out, idx_out = id_index(basin.node_id, id_out)
@@ -742,10 +772,10 @@ function update_jac_prototype!(
     node::Union{Pump, Outlet, TabulatedRatingCurve, User},
 )::Nothing
     (; basin, fractional_flow, connectivity) = p
-    (; graph_flow) = connectivity
+    (; graph) = connectivity
 
     for (i, id) in enumerate(node.node_id)
-        id_in = only(inneighbors(graph_flow, id))
+        id_in = only(inneighbor_labels_type(graph, id, EdgeType.flow))
 
         if hasfield(typeof(node), :is_pid_controlled) && node.is_pid_controlled[i]
             continue
@@ -761,10 +791,10 @@ function update_jac_prototype!(
             jac_prototype[idx_in, idx_in] = 1.0
 
             _, idxs_out =
-                get_fractional_flow_connected_basins(id, basin, fractional_flow, graph_flow)
+                get_fractional_flow_connected_basins(id, basin, fractional_flow, graph)
 
             if isempty(idxs_out)
-                id_out = only(outneighbors(graph_flow, id))
+                id_out = only(outneighbor_labels_type(graph, id, EdgeType.flow))
                 has_index_out, idx_out = id_index(basin.node_id, id_out)
 
                 if has_index_out
@@ -792,7 +822,7 @@ function update_jac_prototype!(
     node::PidControl,
 )::Nothing
     (; basin, connectivity, pump) = p
-    (; graph_control, graph_flow) = connectivity
+    (; graph) = connectivity
 
     n_basins = length(basin.node_id)
 
@@ -801,7 +831,7 @@ function update_jac_prototype!(
         id = node.node_id[i]
 
         # ID of controlled pump/outlet
-        id_controlled = only(outneighbors(graph_control, id))
+        id_controlled = only(outneighbor_labels_type(graph, id, EdgeType.control))
 
         _, listen_idx = id_index(basin.node_id, listen_node_id)
 
@@ -849,17 +879,18 @@ Get the node type specific indices of the fractional flows and basins,
 that are consecutively connected to a node of given id.
 """
 function get_fractional_flow_connected_basins(
-    node_id::Int,
+    node_id::NodeID,
     basin::Basin,
     fractional_flow::FractionalFlow,
-    graph_flow::DiGraph{Int},
+    graph::MetaGraph,
 )::Tuple{Vector{Int}, Vector{Int}}
     fractional_flow_idxs = Int[]
     basin_idxs = Int[]
 
-    for first_outneighbor_id in outneighbors(graph_flow, node_id)
+    for first_outneighbor_id in outneighbor_labels_type(graph, node_id, EdgeType.flow)
         if first_outneighbor_id in fractional_flow.node_id
-            second_outneighbor_id = only(outneighbors(graph_flow, first_outneighbor_id))
+            second_outneighbor_id =
+                only(outneighbor_labels_type(graph, first_outneighbor_id, EdgeType.flow))
             has_index, basin_idx = id_index(basin.node_id, second_outneighbor_id)
             if has_index
                 push!(
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 4cbc2339b..89dba6012 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -406,38 +406,56 @@ function sorted_table!(
     return table
 end
 
+struct NodeID
+    id::Int
+end
+
+function Base.isless(label_1::NodeID, label_2::NodeID)::Bool
+    return label_1.id < label_2.id
+end
+
+function Base.getindex(M::AbstractMatrix{T}, id_row::NodeID, id_col::NodeID)::T where {T}
+    return M[id_row.id, id_col.id]
+end
+
+function Base.setindex!(
+    M::AbstractMatrix{T},
+    value::T,
+    id_row::NodeID,
+    id_col::NodeID,
+)::Nothing where {T}
+    M[id_row.id, id_col.id] = value
+    return nothing
+end
+
+struct EdgeID
+    id::Int
+end
+
 """
 Test for each node given its node type whether the nodes that
 # are downstream ('down-edge') of this node are of an allowed type
 """
-function valid_edges(
-    edge_ids::Dictionary{Tuple{Int, Int}, Int},
-    edge_connection_types::Dictionary{Int, Tuple{Symbol, Symbol}},
-)::Bool
-    rev_edge_ids = dictionary((v => k for (k, v) in pairs(edge_ids)))
-    errors = String[]
-    for (edge_id, (from_type, to_type)) in pairs(edge_connection_types)
-        if !(to_type in neighbortypes(from_type))
-            a, b = rev_edge_ids[edge_id]
-            push!(
-                errors,
-                "Cannot connect a $from_type to a $to_type (edge #$edge_id from node #$a to #$b).",
-            )
+function valid_edges(graph::MetaGraph)::Bool
+    errors = false
+    for e in edges(graph)
+        label_src = label_for(graph, e.src)
+        label_dst = label_for(graph, e.dst)
+        type_src = graph[label_src].type
+        type_dst = graph[label_dst].type
+        if !(type_dst in neighbortypes(type_src))
+            edge_id = graph[label_src, label_dist].id.id
+            @error "Cannot connect a $type_src to a $type_dst (edge #$edge_id from node #$(label_src.id) to #$(label_dst.id))."
         end
     end
-    if isempty(errors)
-        return true
-    else
-        foreach(x -> @error(x), errors)
-        return false
-    end
+    return !errors
 end
 
 """
 Check whether the profile data has no repeats in the levels and the areas start positive.
 """
 function valid_profiles(
-    node_id::Indices{Int},
+    node_id::Indices{NodeID},
     level::Vector{Vector{Float64}},
     area::Vector{Vector{Float64}},
 )::Vector{String}
@@ -469,16 +487,16 @@ end
 Test whether static or discrete controlled flow rates are indeed non-negative.
 """
 function valid_flow_rates(
-    node_id::Vector{Int},
+    node_id::Vector{NodeID},
     flow_rate::Vector,
-    control_mapping::Dict{Tuple{Int, String}, NamedTuple},
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple},
     node_type::Symbol,
 )::Bool
     errors = false
 
     # Collect ids of discrete controlled nodes so that they do not give another error
     # if their initial value is also invalid.
-    ids_controlled = Int[]
+    ids_controlled = NodeID[]
 
     for (key, control_values) in pairs(control_mapping)
         id_controlled = key[1]
@@ -506,12 +524,11 @@ function valid_flow_rates(
 end
 
 function valid_pid_connectivity(
-    pid_control_node_id::Vector{Int},
-    pid_control_listen_node_id::Vector{Int},
-    graph_flow::DiGraph{Int},
-    graph_control::DiGraph{Int},
-    basin_node_id::Indices{Int},
-    pump_node_id::Vector{Int},
+    pid_control_node_id::Vector{NodeID},
+    pid_control_listen_node_id::Vector{NodeID},
+    graph::MetaGraph,
+    basin_node_id::Indices{NodeID},
+    pump_node_id::Vector{NodeID},
 )::Bool
     errors = false
 
@@ -547,24 +564,24 @@ Check that nodes that have fractional flow outneighbors do not have any other ty
 outneighbor, that the fractions leaving a node add up to ≈1 and that the fractions are non-negative.
 """
 function valid_fractional_flow(
-    graph_flow::DiGraph{Int},
-    node_id::Vector{Int},
-    control_mapping::Dict{Tuple{Int64, String}, NamedTuple},
+    graph_flow::MetaGraph,
+    node_id::Vector{NodeID},
+    control_mapping::Dict{Tuple{NodeID, String}, NamedTuple},
 )::Bool
     errors = false
 
     # Node IDs that have fractional flow outneighbors
-    src_ids = Set{Int}()
+    src_ids = Set{NodeID}()
 
     for id in node_id
-        union!(src_ids, inneighbors(graph_flow, id))
+        union!(src_ids, inneighbor_labels(graph_flow, id))
     end
 
-    node_id_set = Set{Int}(node_id)
+    node_id_set = Set{NodeID}(node_id)
     control_states = Set{String}([key[2] for key in keys(control_mapping)])
 
     for src_id in src_ids
-        src_outneighbor_ids = Set(outneighbors(graph_flow, src_id))
+        src_outneighbor_ids = Set(outneighbor_labels(graph_flow, src_id))
         if src_outneighbor_ids ⊈ node_id
             errors = true
             @error(

From aa627a6fb54e5bfaefdc4996e1cd24f1eeff058d Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 15 Nov 2023 18:48:06 +0100
Subject: [PATCH 2/8] the basic model runs without errors

---
 core/src/io.jl         | 33 +++++++++++++----
 core/src/solve.jl      | 81 ++++++++++++++++++++++--------------------
 core/src/utils.jl      | 10 +++---
 core/src/validation.jl |  4 +--
 4 files changed, 75 insertions(+), 53 deletions(-)

diff --git a/core/src/io.jl b/core/src/io.jl
index 66d315cc0..edd630ec1 100644
--- a/core/src/io.jl
+++ b/core/src/io.jl
@@ -148,14 +148,20 @@ end
 "Get the storage and level of all basins as matrices of nbasin × ntime"
 function get_storages_and_levels(
     model::Model,
-)::NamedTuple{
-    (:time, :node_id, :storage, :level),
-    Tuple{Vector{Dates.DateTime}, Vector{Int}, Matrix{Float64}, Matrix{Float64}},
+)::@NamedTuple{
+    time::Vector{Dates.DateTime},
+    node_id::Vector{NodeID},
+    storage::Matrix{Float64},
+    level::Matrix{Float64},
 }
+    NamedTuple{
+        (:time, :node_id, :storage, :level),
+        Tuple{Vector{Dates.DateTime}, Vector{Int}, Matrix{Float64}, Matrix{Float64}},
+    }
     (; config, integrator) = model
     (; sol, p) = integrator
 
-    node_id = p.basin.node_id.values::Vector{Int}
+    node_id = p.basin.node_id.values::Vector{NodeID}
     tsteps = datetime_since.(timesteps(model), config.starttime)
 
     storage = hcat([collect(u_.storage) for u_ in sol.u]...)
@@ -185,15 +191,28 @@ function flow_table(model::Model)::NamedTuple
     (; config, saved_flow, integrator) = model
     (; t, saveval) = saved_flow
     (; connectivity) = integrator.p
+    (; graph) = connectivity
 
-    I, J, _ = findnz(get_tmp(connectivity.flow, integrator.u))
+    I, J, _ = findnz(get_tmp(connectivity.flow, 0))
     # self-loops have no edge ID
-    unique_edge_ids = [get(connectivity.edge_ids_flow, ij, missing) for ij in zip(I, J)]
+    # unique_edge_ids = [get(connectivity.edge_ids_flow, ij, missing) for ij in zip(I, J)]
+    unique_edge_ids_flow = Union{Int, Missing}[]
+    for (i, j) in zip(I, J)
+        if i == j
+            push!(unique_edge_ids_flow, missing)
+        else
+            label_src = label_for(graph, i)
+            label_dst = label_for(graph, j)
+            edge_metadata = graph[label_src, label_dst]
+            push!(unique_edge_ids_flow, edge_metadata.id.id)
+        end
+    end
+
     nflow = length(I)
     ntsteps = length(t)
 
     time = repeat(datetime_since.(t, config.starttime); inner = nflow)
-    edge_id = repeat(unique_edge_ids; outer = ntsteps)
+    edge_id = repeat(unique_edge_ids_flow; outer = ntsteps)
     from_node_id = repeat(I; outer = ntsteps)
     to_node_id = repeat(J; outer = ntsteps)
     flow = FlatVector(saveval)
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 7fd90bbc9..9ae388692 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -601,7 +601,7 @@ function continuous_control!(
     max_flow_rate_pump = pump.max_flow_rate
     min_flow_rate_outlet = outlet.min_flow_rate
     max_flow_rate_outlet = outlet.max_flow_rate
-    (; graph_control, graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, active, target, pid_params, listen_node_id, error) = pid_control
     (; current_area) = basin
 
@@ -626,13 +626,14 @@ function continuous_control!(
         listened_node_id = listen_node_id[i]
         _, listened_node_idx = id_index(basin.node_id, listened_node_id)
 
-        controlled_node_id = only(outneighbors(graph_control, id))
+        controlled_node_id =
+            only(outneighbor_labels_type(graph_control, id, EdgeType.control))
         controls_pump = (controlled_node_id in pump.node_id)
 
         # No flow of outlet if source level is lower than target level
         if !controls_pump
-            src_id = only(inneighbors(graph_flow, controlled_node_id))
-            dst_id = only(outneighbors(graph_flow, controlled_node_id))
+            src_id = only(inneighbor_labels_type(graph, controlled_node_id, EdgeType.flow))
+            dst_id = only(outneighbor_labels_type(graph, controlled_node_id, EdgeType.flow))
 
             src_level = get_level(p, src_id, t; storage)
             dst_level = get_level(p, dst_id, t; storage)
@@ -656,7 +657,8 @@ function continuous_control!(
             controlled_node_idx = findsorted(outlet.node_id, controlled_node_id)
 
             # Upstream node of outlet does not have to be a basin
-            upstream_node_id = only(inneighbors(graph_flow, controlled_node_id))
+            upstream_node_id =
+                only(inneighbor_labels_type(graph, controlled_node_id, EdgeType.flow))
             has_index, upstream_basin_idx = id_index(basin.node_id, upstream_node_id)
             if has_index
                 upstream_basin_storage = u.storage[upstream_basin_idx]
@@ -723,8 +725,8 @@ function continuous_control!(
         end
 
         # Set flow for connected edges
-        src_id = only(inneighbors(graph_flow, controlled_node_id))
-        dst_id = only(outneighbors(graph_flow, controlled_node_id))
+        src_id = only(inneighbor_labels_type(graph, controlled_node_id, EdgeType.flow))
+        dst_id = only(outneighbor_labels_type(graph, controlled_node_id, EdgeType.flow))
 
         flow[src_id, controlled_node_id] = flow_rate
         flow[controlled_node_id, dst_id] = flow_rate
@@ -738,7 +740,8 @@ function continuous_control!(
         if controls_pump
             for id in outneighbors(graph_flow, controlled_node_id)
                 if id in fractional_flow.node_id
-                    after_ff_id = only(outneighbours(graph_flow, id))
+                    after_ff_id =
+                        only(outneighbor_labels_type(graph_flow, id, EdgeType.flow))
                     ff_idx = findsorted(fractional_flow, id)
                     flow_rate_fraction = fractional_flow.fraction[ff_idx] * flow_rate
                     flow[id, after_ff_id] = flow_rate_fraction
@@ -762,14 +765,14 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity, basin) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, allocated, demand, active, return_factor, min_level) = user
 
     flow = get_tmp(flow, storage)
 
     for (i, id) in enumerate(node_id)
-        src_id = only(inneighbors(graph_flow, id))
-        dst_id = only(outneighbors(graph_flow, id))
+        src_id = only(inneighbor_labels_type(graph_flow, id, EdgeType.flow))
+        dst_id = only(outneighbor_labels_type(graph_flow, id, EdgeType.flow))
 
         if !active[i]
             continue
@@ -817,12 +820,12 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, active, resistance) = linear_resistance
     flow = get_tmp(flow, storage)
     for (i, id) in enumerate(node_id)
-        basin_a_id = only(inneighbors(graph_flow, id))
-        basin_b_id = only(outneighbors(graph_flow, id))
+        basin_a_id = only(inneighbor_labels_type(graph, id, EdgeType.flow))
+        basin_b_id = only(outneighbor_labels_type(graph, id, EdgeType.flow))
 
         if active[i]
             q =
@@ -847,12 +850,12 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; basin, connectivity) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, active, tables) = tabulated_rating_curve
     flow = get_tmp(flow, storage)
     for (i, id) in enumerate(node_id)
-        upstream_basin_id = only(inneighbors(graph_flow, id))
-        downstream_ids = outneighbors(graph_flow, id)
+        upstream_basin_id = only(inneighbor_labels_type(graph, id, EdgeType.flow))
+        downstream_ids = outneighbor_labels_type(graph, id, EdgeType.flow)
 
         if active[i]
             hasindex, basin_idx = id_index(basin.node_id, upstream_basin_id)
@@ -917,13 +920,13 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; basin, connectivity) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, active, length, manning_n, profile_width, profile_slope) =
         manning_resistance
     flow = get_tmp(flow, storage)
     for (i, id) in enumerate(node_id)
-        basin_a_id = only(inneighbors(graph_flow, id))
-        basin_b_id = only(outneighbors(graph_flow, id))
+        basin_a_id = only(inneighbor_labels_type(graph, id, EdgeType.flow))
+        basin_b_id = only(outneighbor_labels_type(graph, id, EdgeType.flow))
 
         if !active[i]
             continue
@@ -974,13 +977,13 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, fraction) = fractional_flow
     flow = get_tmp(flow, storage)
 
     for (i, id) in enumerate(node_id)
-        downstream_id = only(outneighbors(graph_flow, id))
-        upstream_id = only(inneighbors(graph_flow, id))
+        downstream_id = only(outneighbor_labels_type(graph, id, EdgeType.flow))
+        upstream_id = only(inneighbor_labels_type(graph, id, EdgeType.flow))
         # overwrite the inflow such that flow is conserved over the FractionalFlow
         outflow = flow[upstream_id, id] * fraction[i]
         flow[upstream_id, id] = outflow
@@ -996,12 +999,12 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id) = terminal
     flow = get_tmp(flow, storage)
 
     for id in node_id
-        for upstream_id in inneighbors(graph_flow, id)
+        for upstream_id in inneighbor_labels_type(graph, id, EdgeType.flow)
             q = flow[upstream_id, id]
             flow[id, id] -= q
         end
@@ -1016,16 +1019,16 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id) = level_boundary
     flow = get_tmp(flow, storage)
 
     for id in node_id
-        for in_id in inneighbors(graph_flow, id)
+        for in_id in inneighbor_labels_type(graph, id, EdgeType.flow)
             q = flow[in_id, id]
             flow[id, id] -= q
         end
-        for out_id in outneighbors(graph_flow, id)
+        for out_id in outneighbor_labels_type(graph, id, EdgeType.flow)
             q = flow[id, out_id]
             flow[id, id] += q
         end
@@ -1040,13 +1043,13 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, active, flow_rate) = flow_boundary
     flow = get_tmp(flow, storage)
 
     for (i, id) in enumerate(node_id)
         # Requirement: edge points away from the flow boundary
-        for dst_id in outneighbors(graph_flow, id)
+        for dst_id in outneighbor_labels_type(graph, id, EdgeType.flow)
             if !active[i]
                 continue
             end
@@ -1067,14 +1070,14 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity, basin) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, active, flow_rate, is_pid_controlled) = pump
     flow = get_tmp(flow, storage)
     flow_rate = get_tmp(flow_rate, storage)
     for (id, isactive, rate, pid_controlled) in
         zip(node_id, active, flow_rate, is_pid_controlled)
-        src_id = only(inneighbors(graph_flow, id))
-        dst_id = only(outneighbors(graph_flow, id))
+        src_id = only(inneighbor_labels_type(graph, id, EdgeType.flow))
+        dst_id = only(outneighbor_labels_type(graph, id, EdgeType.flow))
 
         if !isactive || pid_controlled
             continue
@@ -1102,13 +1105,13 @@ function formulate_flow!(
     t::Float64,
 )::Nothing
     (; connectivity, basin) = p
-    (; graph_flow, flow) = connectivity
+    (; graph, flow) = connectivity
     (; node_id, active, flow_rate, is_pid_controlled, min_crest_level) = outlet
     flow = get_tmp(flow, storage)
     flow_rate = get_tmp(flow_rate, storage)
     for (i, id) in enumerate(node_id)
-        src_id = only(inneighbors(graph_flow, id))
-        dst_id = only(outneighbors(graph_flow, id))
+        src_id = only(inneighbor_labels_type(graph_flow, id, EdgeType.flow))
+        dst_id = only(outneighbor_labels_type(graph_flow, id, EdgeType.flow))
 
         if !active[i] || is_pid_controlled[i]
             continue
@@ -1152,12 +1155,12 @@ function formulate_du!(
     # loop over basins
     # subtract all outgoing flows
     # add all ingoing flows
-    (; graph_flow) = connectivity
+    (; graph) = connectivity
     for (i, basin_id) in enumerate(basin.node_id)
-        for in_id in inneighbors(graph_flow, basin_id)
+        for in_id in inneighbor_labels_type(graph, basin_id, EdgeType.flow)
             du[i] += flow[in_id, basin_id]
         end
-        for out_id in outneighbors(graph_flow, basin_id)
+        for out_id in outneighbor_labels_type(graph, basin_id, EdgeType.flow)
             du[i] -= flow[basin_id, out_id]
         end
     end
diff --git a/core/src/utils.jl b/core/src/utils.jl
index 5c97ce18f..582f9fe4e 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -439,7 +439,7 @@ storage: tells ForwardDiff whether this call is for differentiation or not
 """
 function get_level(
     p::Parameters,
-    node_id::Int,
+    node_id::NodeID,
     t::Float64;
     storage::Union{AbstractArray, Number} = 0,
 )::Union{Real, Nothing}
@@ -467,7 +467,7 @@ function id_index(ids::Indices{NodeID}, id::NodeID)::Tuple{Bool, Int}
 end
 
 "Return the bottom elevation of the basin with index i, or nothing if it doesn't exist"
-function basin_bottom(basin::Basin, node_id::Int)::Union{Float64, Nothing}
+function basin_bottom(basin::Basin, node_id::NodeID)::Union{Float64, Nothing}
     hasindex, i = id_index(basin.node_id, node_id)
     return if hasindex
         # get level(storage) interpolation function
@@ -482,9 +482,9 @@ end
 "Get the bottom on both ends of a node. If only one has a bottom, use that for both."
 function basin_bottoms(
     basin::Basin,
-    basin_a_id::Int,
-    basin_b_id::Int,
-    id::Int,
+    basin_a_id::NodeID,
+    basin_b_id::NodeID,
+    id::NodeID,
 )::Tuple{Float64, Float64}
     bottom_a = basin_bottom(basin, basin_a_id)
     bottom_b = basin_bottom(basin, basin_b_id)
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 89dba6012..943924f42 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -414,12 +414,12 @@ function Base.isless(label_1::NodeID, label_2::NodeID)::Bool
     return label_1.id < label_2.id
 end
 
-function Base.getindex(M::AbstractMatrix{T}, id_row::NodeID, id_col::NodeID)::T where {T}
+function Base.getindex(M::AbstractArray, id_row::NodeID, id_col::NodeID)
     return M[id_row.id, id_col.id]
 end
 
 function Base.setindex!(
-    M::AbstractMatrix{T},
+    M::AbstractArray,
     value::T,
     id_row::NodeID,
     id_col::NodeID,

From 4b3e90912b5b6a9375ca2010c022dc9548f6f5f4 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 15 Nov 2023 19:21:53 +0100
Subject: [PATCH 3/8] Get metadata from edge of codes

---
 core/src/io.jl    | 4 +---
 core/src/utils.jl | 6 ++++++
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/core/src/io.jl b/core/src/io.jl
index edd630ec1..3f058612b 100644
--- a/core/src/io.jl
+++ b/core/src/io.jl
@@ -201,9 +201,7 @@ function flow_table(model::Model)::NamedTuple
         if i == j
             push!(unique_edge_ids_flow, missing)
         else
-            label_src = label_for(graph, i)
-            label_dst = label_for(graph, j)
-            edge_metadata = graph[label_src, label_dst]
+            edge_metadata = metadata_from_edge(graph, Edge(i, j))
             push!(unique_edge_ids_flow, edge_metadata.id.id)
         end
     end
diff --git a/core/src/utils.jl b/core/src/utils.jl
index 582f9fe4e..577621be7 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -63,6 +63,12 @@ function outneighbor_labels_type(
     ]
 end
 
+function metadata_from_edge(graph::MetaGraph, edge::Edge{Int})::EdgeMetadata
+    label_src = label_for(graph, edge.src)
+    label_dst = label_for(graph, edge.dst)
+    return graph[label_src, label_dst]
+end
+
 "Calculate a profile storage by integrating the areas over the levels"
 function profile_storage(levels::Vector, areas::Vector)::Vector{Float64}
     # profile starts at the bottom; first storage is 0

From 5caae000d8a64e25aa7b9ea2cb818d1af104d90c Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 16 Nov 2023 15:59:07 +0100
Subject: [PATCH 4/8] More models run

---
 core/src/allocation.jl | 17 +++++++++--------
 core/src/bmi.jl        | 39 +++++++++++++++++++++------------------
 core/src/create.jl     | 27 ++++++++++++++-------------
 core/src/io.jl         |  2 +-
 core/src/solve.jl      | 21 ++++++++++-----------
 core/src/utils.jl      | 37 +++++++++++++++++++++++++------------
 core/src/validation.jl | 35 ++++++++++++++++++++---------------
 7 files changed, 100 insertions(+), 78 deletions(-)

diff --git a/core/src/allocation.jl b/core/src/allocation.jl
index c893cf4aa..aad059a53 100644
--- a/core/src/allocation.jl
+++ b/core/src/allocation.jl
@@ -5,26 +5,27 @@ Get:
 """
 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
@@ -303,7 +304,7 @@ 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
@@ -642,7 +643,7 @@ An AllocationModel object.
 function AllocationModel(
     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
diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index 26c86bbe4..fd7759b99 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -253,15 +253,15 @@ function discrete_control_condition(out, u, t, integrator)
     (; p) = integrator
     (; discrete_control) = p
 
-    for (i, (listen_feature_id, variable, greater_than, look_ahead)) in enumerate(
+    for (i, (listen_node_id, variable, greater_than, look_ahead)) in enumerate(
         zip(
-            discrete_control.listen_feature_id,
+            discrete_control.listen_node_id,
             discrete_control.variable,
             discrete_control.greater_than,
             discrete_control.look_ahead,
         ),
     )
-        value = get_value(p, listen_feature_id, variable, look_ahead, u, t)
+        value = get_value(p, listen_node_id, variable, look_ahead, u, t)
         diff = value - greater_than
         out[i] = diff
     end
@@ -273,7 +273,7 @@ from flow boundaries.
 """
 function get_value(
     p::Parameters,
-    feature_id::Int,
+    node_id::NodeID,
     variable::String,
     Δt::Float64,
     u::AbstractVector{Float64},
@@ -282,8 +282,8 @@ function get_value(
     (; basin, flow_boundary, level_boundary) = p
 
     if variable == "level"
-        hasindex_basin, basin_idx = id_index(basin.node_id, feature_id)
-        level_boundary_idx = findsorted(level_boundary.node_id, feature_id)
+        hasindex_basin, basin_idx = id_index(basin.node_id, node_id)
+        level_boundary_idx = findsorted(level_boundary.node_id, node_id)
 
         if hasindex_basin
             _, level = get_area_and_level(basin, basin_idx, u[basin_idx])
@@ -298,7 +298,7 @@ function get_value(
         value = level
 
     elseif variable == "flow_rate"
-        flow_boundary_idx = findsorted(flow_boundary.node_id, feature_id)
+        flow_boundary_idx = findsorted(flow_boundary.node_id, node_id)
 
         if flow_boundary_idx === nothing
             error("Flow condition node #$feature_id is not a flow boundary.")
@@ -318,7 +318,7 @@ An upcrossing means that a condition (always greater than) becomes true.
 function discrete_control_affect_upcrossing!(integrator, condition_idx)
     (; p, u, t) = integrator
     (; discrete_control, basin) = p
-    (; variable, condition_value, listen_feature_id) = discrete_control
+    (; variable, condition_value, listen_node_id) = discrete_control
 
     condition_value[condition_idx] = true
 
@@ -330,7 +330,7 @@ function discrete_control_affect_upcrossing!(integrator, condition_idx)
     # only possibly the du. Parameter changes can change the flow on an edge discontinuously,
     # giving the possibility of logical paradoxes where certain parameter changes immediately
     # undo the truth state that caused that parameter change.
-    is_basin = id_index(basin.node_id, discrete_control.listen_feature_id[condition_idx])[1]
+    is_basin = id_index(basin.node_id, discrete_control.listen_node_id[condition_idx])[1]
     # NOTE: The above no longer works when listen feature ids can be something other than node ids
     # I think the more durable option is to give all possible condition types a different variable string,
     # e.g. basin.level and level_boundary.level
@@ -339,7 +339,7 @@ function discrete_control_affect_upcrossing!(integrator, condition_idx)
         # du for the basin of this level condition
         du = zero(u)
         water_balance!(du, u, p, t)
-        _, condition_basin_idx = id_index(basin.node_id, listen_feature_id[condition_idx])
+        _, condition_basin_idx = id_index(basin.node_id, listen_node_id[condition_idx])
 
         if du[condition_basin_idx] < 0.0
             condition_value[condition_idx] = false
@@ -354,7 +354,7 @@ An downcrossing means that a condition (always greater than) becomes false.
 function discrete_control_affect_downcrossing!(integrator, condition_idx)
     (; p, u, t) = integrator
     (; discrete_control, basin) = p
-    (; variable, condition_value, listen_feature_id) = discrete_control
+    (; variable, condition_value, listen_node_id) = discrete_control
 
     condition_value[condition_idx] = false
 
@@ -372,7 +372,7 @@ function discrete_control_affect_downcrossing!(integrator, condition_idx)
         du = zero(u)
         water_balance!(du, u, p, t)
         has_index, condition_basin_idx =
-            id_index(basin.node_id, listen_feature_id[condition_idx])
+            id_index(basin.node_id, listen_node_id[condition_idx])
 
         if has_index && du[condition_basin_idx] > 0.0
             condition_value[condition_idx] = true
@@ -445,13 +445,16 @@ function discrete_control_affect!(
         record = discrete_control.record
 
         push!(record.time, integrator.t)
-        push!(record.control_node_id, discrete_control_node_id)
+        push!(record.control_node_id, discrete_control_node_id.value)
         push!(record.truth_state, truth_state_used)
         push!(record.control_state, control_state_new)
 
         # Loop over nodes which are under control of this control node
-        for target_node_id in
-            outneighbors(connectivity.graph_control, discrete_control_node_id)
+        for target_node_id in outneighbor_labels_type(
+            connectivity.graph,
+            discrete_control_node_id,
+            EdgeType.control,
+        )
             set_control_params!(p, target_node_id, control_state_new)
         end
 
@@ -461,8 +464,8 @@ function discrete_control_affect!(
     return control_state_change
 end
 
-function set_control_params!(p::Parameters, node_id::Int, control_state::String)
-    node = getfield(p, p.lookup[node_id])
+function set_control_params!(p::Parameters, node_id::NodeID, control_state::String)
+    node = getfield(p, p.connectivity.graph[node_id].type)
     idx = searchsortedfirst(node.node_id, node_id)
     new_state = node.control_mapping[(node_id, control_state)]
 
@@ -496,7 +499,7 @@ function update_basin(integrator)::Nothing
     )
 
     for row in timeblock
-        hasindex, i = id_index(node_id, row.node_id)
+        hasindex, i = id_index(node_id, NodeID(row.node_id))
         @assert hasindex "Table 'Basin / time' contains non-Basin IDs"
         set_table_row!(table, row, i)
     end
diff --git a/core/src/create.jl b/core/src/create.jl
index f1eba3726..a622885dc 100644
--- a/core/src/create.jl
+++ b/core/src/create.jl
@@ -206,11 +206,11 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity
     # Build the sparsity structure of the flow matrix
     flow = spzeros(nv(graph), nv(graph))
     for e in edges(graph)
-        label_src = label_for(graph, e.src)
-        label_dst = label_for(graph, e.dst)
-        edge_metadata = graph[label_src, label_dst]
+        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
-            flow[label_src.id, label_dst.id] = 1.0
+            flow[id_src, id_dst] = 1.0
         end
     end
 
@@ -275,7 +275,7 @@ function generate_allocation_models!(p::Parameters, db::DB, config::Config)::Not
             AllocationModel(
                 allocation_network_id,
                 p,
-                allocation_group_node.fid,
+                NodeID.(allocation_group_node.fid),
                 source_edge_ids,
                 config.allocation.timestep,
             ),
@@ -590,22 +590,22 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
     condition = load_structvector(db, config, DiscreteControlConditionV1)
 
     condition_value = fill(false, length(condition.node_id))
-    control_state::Dict{Int, Tuple{String, Float64}} = Dict()
+    control_state::Dict{NodeID, Tuple{String, Float64}} = Dict()
 
     rows = execute(db, "select from_node_id, edge_type from Edge")
     for (; from_node_id, edge_type) in rows
         if edge_type == "control"
-            control_state[from_node_id] = ("undefined_state", 0.0)
+            control_state[NodeID(from_node_id)] = ("undefined_state", 0.0)
         end
     end
 
     logic = load_structvector(db, config, DiscreteControlLogicV1)
 
-    logic_mapping = Dict{Tuple{Int, String}, String}()
+    logic_mapping = Dict{Tuple{NodeID, String}, String}()
 
     for (node_id, truth_state, control_state_) in
         zip(logic.node_id, logic.truth_state, logic.control_state)
-        logic_mapping[(node_id, truth_state)] = control_state_
+        logic_mapping[(NodeID(node_id), truth_state)] = control_state_
     end
 
     logic_mapping = expand_logic_mapping(logic_mapping)
@@ -619,8 +619,8 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
     )
 
     return DiscreteControl(
-        condition.node_id, # Not unique
-        condition.listen_feature_id,
+        NodeID.(condition.node_id), # Not unique
+        NodeID.(condition.listen_feature_id),
         condition.variable,
         look_ahead,
         condition.greater_than,
@@ -797,7 +797,7 @@ function User(db::DB, config::Config)::User
     )
 
     return User(
-        node_ids,
+        NodeID.(node_ids),
         active,
         interpolations,
         allocated,
@@ -831,7 +831,8 @@ function Parameters(db::DB, config::Config)::Parameters
 
     # Set is_pid_controlled to true for those pumps and outlets that are PID controlled
     for id in pid_control.node_id
-        id_controlled = only(outneighbors(connectivity.graph_control, id))
+        id_controlled =
+            only(outneighbor_labels_type(connectivity.graph, id, EdgeType.control))
         pump_idx = findsorted(pump.node_id, id_controlled)
         if pump_idx === nothing
             outlet_idx = findsorted(outlet.node_id, id_controlled)
diff --git a/core/src/io.jl b/core/src/io.jl
index 3f058612b..8cf74318f 100644
--- a/core/src/io.jl
+++ b/core/src/io.jl
@@ -202,7 +202,7 @@ function flow_table(model::Model)::NamedTuple
             push!(unique_edge_ids_flow, missing)
         else
             edge_metadata = metadata_from_edge(graph, Edge(i, j))
-            push!(unique_edge_ids_flow, edge_metadata.id.id)
+            push!(unique_edge_ids_flow, edge_metadata.id.value)
         end
     end
 
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 9ae388692..e7aa3acd9 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -21,9 +21,9 @@ problem: The JuMP.jl model for solving the allocation problem
 """
 struct AllocationModel
     allocation_network_id::Int
-    node_id::Vector{Int}
-    node_id_mapping::Dict{Int, Tuple{Int, Symbol}}
-    node_id_mapping_inverse::Dict{Int, Tuple{Int, Symbol}}
+    node_id::Vector{NodeID}
+    node_id_mapping::Dict{NodeID, Tuple{Int, Symbol}}
+    node_id_mapping_inverse::Dict{Int, Tuple{NodeID, Symbol}}
     allocgraph_edge_ids_user_demand::Dict{Int, Int}
     source_edge_mapping::Dict{Int, Int}
     graph_allocation::DiGraph{Int}
@@ -384,13 +384,13 @@ record: Namedtuple with discrete control information for results
 """
 struct DiscreteControl <: AbstractParameterNode
     node_id::Vector{NodeID}
-    listen_feature_id::Vector{Int}
+    listen_node_id::Vector{NodeID}
     variable::Vector{String}
     look_ahead::Vector{Float64}
     greater_than::Vector{Float64}
     condition_value::Vector{Bool}
-    control_state::Dict{Int, Tuple{String, Float64}}
-    logic_mapping::Dict{Tuple{Int, String}, String}
+    control_state::Dict{NodeID, Tuple{String, Float64}}
+    logic_mapping::Dict{Tuple{NodeID, String}, String}
     record::@NamedTuple{
         time::Vector{Float64},
         control_node_id::Vector{Int},
@@ -626,8 +626,7 @@ function continuous_control!(
         listened_node_id = listen_node_id[i]
         _, listened_node_idx = id_index(basin.node_id, listened_node_id)
 
-        controlled_node_id =
-            only(outneighbor_labels_type(graph_control, id, EdgeType.control))
+        controlled_node_id = only(outneighbor_labels_type(graph, id, EdgeType.control))
         controls_pump = (controlled_node_id in pump.node_id)
 
         # No flow of outlet if source level is lower than target level
@@ -738,7 +737,7 @@ function continuous_control!(
 
         # When the controlled pump flows out into fractional flow nodes
         if controls_pump
-            for id in outneighbors(graph_flow, controlled_node_id)
+            for id in outneighbor_labels_type(graph, controlled_node_id, EdgeType.flow)
                 if id in fractional_flow.node_id
                     after_ff_id =
                         only(outneighbor_labels_type(graph_flow, id, EdgeType.flow))
@@ -1110,8 +1109,8 @@ function formulate_flow!(
     flow = get_tmp(flow, storage)
     flow_rate = get_tmp(flow_rate, storage)
     for (i, id) in enumerate(node_id)
-        src_id = only(inneighbor_labels_type(graph_flow, id, EdgeType.flow))
-        dst_id = only(outneighbor_labels_type(graph_flow, id, EdgeType.flow))
+        src_id = only(inneighbor_labels_type(graph, id, EdgeType.flow))
+        dst_id = only(outneighbor_labels_type(graph, id, EdgeType.flow))
 
         if !active[i] || is_pid_controlled[i]
             continue
diff --git a/core/src/utils.jl b/core/src/utils.jl
index 577621be7..edd46d6d6 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -25,7 +25,8 @@ function create_graph(db::DB)::MetaGraph
     for (i, row) in enumerate(node_rows)
         allocation_network_id =
             hasproperty(row, :allocation_network_id) ? row.allocation_network_id : 0
-        graph[NodeID(row.fid)] = NodeMetadata(Symbol(row.type), i, allocation_network_id)
+        graph[NodeID(row.fid)] =
+            NodeMetadata(Symbol(snake_case(row.type)), i, allocation_network_id)
     end
     for (; from_node_id, to_node_id, edge_type, fid) in edge_rows
         if edge_type == "flow"
@@ -63,6 +64,17 @@ function outneighbor_labels_type(
     ]
 end
 
+function all_neighbor_labels_type(
+    graph::MetaGraph,
+    label::NodeID,
+    edge_type::EdgeType.T,
+)::Vector{NodeID}
+    return [
+        outneighbor_labels_type(graph, label, edge_type)...,
+        inneighbor_labels_type(graph, label, edge_type)...,
+    ]
+end
+
 function metadata_from_edge(graph::MetaGraph, edge::Edge{Int})::EdgeMetadata
     label_src = label_for(graph, edge.src)
     label_dst = label_for(graph, edge.dst)
@@ -526,7 +538,7 @@ Check:
 function valid_discrete_control(p::Parameters, config::Config)::Bool
     (; discrete_control, connectivity) = p
     (; graph) = connectivity
-    (; node_id, logic_mapping, look_ahead, variable, listen_feature_id) = discrete_control
+    (; node_id, logic_mapping, look_ahead, variable, listen_node_id) = discrete_control
 
     t_end = seconds_since(config.endtime, config.starttime)
     errors = false
@@ -560,10 +572,10 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
 
         # Check whether these control states are defined for the
         # control outneighbors
-        for id_outneighbor in outneighbors(graph_control, id)
+        for id_outneighbor in outneighbor_labels_type(graph, id, EdgeType.control)
 
             # Node object for the outneighbor node type
-            node = getfield(p, lookup[id_outneighbor])
+            node = getfield(p, graph[id_outneighbor].type)
 
             # Get control states of the controlled node
             control_states_controlled = Set{String}()
@@ -587,9 +599,9 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
             end
         end
     end
-    for (Δt, var, feature_id) in zip(look_ahead, variable, listen_feature_id)
+    for (Δt, var, node_id) in zip(look_ahead, variable, listen_node_id)
         if !iszero(Δt)
-            node_type = p.lookup[feature_id]
+            node_type = p.connectivity.graph[node_id].type
             # TODO: If more transient listen variables must be supported, this validation must be more specific
             # (e.g. for some node some variables are transient, some not).
             if node_type ∉ [:flow_boundary, :level_boundary]
@@ -604,7 +616,7 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
                     idx = if node_type == :Basin
                         id_index(node.node_id, feature_id)
                     else
-                        searchsortedfirst(node.node_id, feature_id)
+                        searchsortedfirst(node.node_id, node_id)
                     end
                     interpolation = getfield(node, Symbol(var))[idx]
                     if t_end + Δt > interpolation.t[end]
@@ -623,9 +635,9 @@ Replace the truth states in the logic mapping which contain wildcards with
 all possible explicit truth states.
 """
 function expand_logic_mapping(
-    logic_mapping::Dict{Tuple{Int, String}, String},
-)::Dict{Tuple{Int, String}, String}
-    logic_mapping_expanded = Dict{Tuple{Int, String}, String}()
+    logic_mapping::Dict{Tuple{NodeID, String}, String},
+)::Dict{Tuple{NodeID, String}, String}
+    logic_mapping_expanded = Dict{Tuple{NodeID, String}, String}()
 
     for (node_id, truth_state) in keys(logic_mapping)
         pattern = r"^[TFUD\*]+$"
@@ -850,7 +862,7 @@ function update_jac_prototype!(
         jac_prototype[pid_state_idx, listen_idx] = 1.0
 
         if id_controlled in pump.node_id
-            id_pump_out = only(inneighbors(graph_flow, id_controlled))
+            id_pump_out = only(inneighbor_labels_type(graph, id_controlled, EdgeType.flow))
 
             # The basin downstream of the pump
             has_index, idx_out_out = id_index(basin.node_id, id_pump_out)
@@ -863,7 +875,8 @@ function update_jac_prototype!(
                 jac_prototype[listen_idx, idx_out_out] = 1.0
             end
         else
-            id_outlet_in = only(outneighbors(graph_flow, id_controlled))
+            id_outlet_in =
+                only(outneighbor_labels_type(graph, id_controlled, EdgeType.flow))
 
             # The basin upstream of the outlet
             has_index, idx_out_in = id_index(basin.node_id, id_outlet_in)
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 943924f42..b3a53cf6a 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -407,15 +407,18 @@ function sorted_table!(
 end
 
 struct NodeID
-    id::Int
+    value::Int
 end
 
-function Base.isless(label_1::NodeID, label_2::NodeID)::Bool
-    return label_1.id < label_2.id
+Base.broadcastable(id::NodeID) = Ref(id)
+Base.show(io::IO, id::NodeID) = print(io, "#$(id.value)")
+
+function Base.isless(id_1::NodeID, id_2::NodeID)::Bool
+    return id_1.value < id_2.value
 end
 
 function Base.getindex(M::AbstractArray, id_row::NodeID, id_col::NodeID)
-    return M[id_row.id, id_col.id]
+    return M[id_row.value, id_col.value]
 end
 
 function Base.setindex!(
@@ -424,12 +427,12 @@ function Base.setindex!(
     id_row::NodeID,
     id_col::NodeID,
 )::Nothing where {T}
-    M[id_row.id, id_col.id] = value
+    M[id_row.value, id_col.value] = value
     return nothing
 end
 
 struct EdgeID
-    id::Int
+    value::Int
 end
 
 """
@@ -439,13 +442,13 @@ Test for each node given its node type whether the nodes that
 function valid_edges(graph::MetaGraph)::Bool
     errors = false
     for e in edges(graph)
-        label_src = label_for(graph, e.src)
-        label_dst = label_for(graph, e.dst)
-        type_src = graph[label_src].type
-        type_dst = graph[label_dst].type
+        id_src = label_for(graph, e.src)
+        id_dst = label_for(graph, e.dst)
+        type_src = graph[id_src].type
+        type_dst = graph[id_dst].type
         if !(type_dst in neighbortypes(type_src))
-            edge_id = graph[label_src, label_dist].id.id
-            @error "Cannot connect a $type_src to a $type_dst (edge #$edge_id from node #$(label_src.id) to #$(label_dst.id))."
+            edge_id = graph[id_src, id_dst].id.value
+            @error "Cannot connect a $type_src to a $type_dst (edge #$edge_id from node #$id_src to #$id_dst)."
         end
     end
     return !errors
@@ -539,16 +542,18 @@ function valid_pid_connectivity(
             errors = true
         end
 
-        controlled_id = only(outneighbors(graph_control, id))
+        controlled_id = only(outneighbor_labels_type(graph, id, EdgeType.control))
 
         if controlled_id in pump_node_id
-            pump_intake_id = only(inneighbors(graph_flow, controlled_id))
+            pump_intake_id =
+                only(inneighbor_labels_type(graph, controlled_id, EdgeType.flow))
             if pump_intake_id != listen_id
                 @error "Listen node #$listen_id of PidControl node #$id is not upstream of controlled pump #$controlled_id"
                 errors = true
             end
         else
-            outlet_outflow_id = only(outneighbors(graph_flow, controlled_id))
+            outlet_outflow_id =
+                only(outneighbor_labels_type(graph, controlled_id, EdgeType.flow))
             if outlet_outflow_id != listen_id
                 @error "Listen node #$listen_id of PidControl node #$id is not downstream of controlled outlet #$controlled_id"
                 errors = true

From 50af899507753e577ece0f06c749c74f18ddc861 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Fri, 17 Nov 2023 13:34:31 +0100
Subject: [PATCH 5/8] All tests are green...?

---
 core/Project.toml       |   3 ++
 core/src/allocation.jl  |  84 ++++++++++++++++++++-----------
 core/src/bmi.jl         |   6 +--
 core/src/solve.jl       |  15 +++---
 core/src/utils.jl       |  22 ++++----
 core/src/validation.jl  |  80 +++++++++++++++--------------
 core/test/control.jl    |  25 ++++-----
 core/test/run_models.jl |   4 +-
 core/test/utils.jl      |  69 +++++++++++++++----------
 core/test/validation.jl | 109 +++++++++++++++++++++++++++-------------
 10 files changed, 252 insertions(+), 165 deletions(-)

diff --git a/core/Project.toml b/core/Project.toml
index 630c40fda..678c9139e 100644
--- a/core/Project.toml
+++ b/core/Project.toml
@@ -54,14 +54,17 @@ DataStructures = "0.18"
 Dates = "1"
 Dictionaries = "0.3.25"
 DiffEqCallbacks = "2.29.1"
+EnumX = "1.0"
 FiniteDiff = "2.21"
 ForwardDiff = "0.10"
+Graphs = "1.9"
 HiGHS = "1.7"
 IterTools = "1.4"
 JuMP = "1.15"
 Legolas = "0.5"
 Logging = "1"
 LoggingExtras = "1"
+MetaGraphsNext = "0.6"
 OrdinaryDiffEq = "6.7"
 PreallocationTools = "0.4"
 SQLite = "1.5.1"
diff --git a/core/src/allocation.jl b/core/src/allocation.jl
index aad059a53..36a6f53c0 100644
--- a/core/src/allocation.jl
+++ b/core/src/allocation.jl
@@ -1,3 +1,19 @@
+"""
+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
@@ -38,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
@@ -72,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
@@ -153,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]
@@ -182,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]
@@ -208,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
@@ -221,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
@@ -250,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
 
@@ -275,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
@@ -312,7 +333,7 @@ function allocation_graph(
         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
@@ -394,7 +415,7 @@ Non-negativivity constraints are also immediately added to the basin allocation
 """
 function add_variables_allocation_basin!(
     problem::JuMP.Model,
-    node_id_mapping::Dict{Int, Tuple{Int, Symbol}},
+    node_id_mapping::Dict{NodeID, Tuple{Int, Symbol}},
     allocgraph_node_ids_basin::Vector{Int},
 )::Nothing
     JuMP.@variable(problem, A_basin[i = allocgraph_node_ids_basin] >= 0.0)
@@ -571,7 +592,7 @@ end
 Construct the allocation problem for the current subnetwork as a JuMP.jl model.
 """
 function allocation_problem(
-    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},
@@ -718,7 +739,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)
@@ -731,7 +752,7 @@ 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)
@@ -739,7 +760,10 @@ function assign_allocations!(
         # 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
@@ -750,14 +774,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
diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index fd7759b99..4a7fcda34 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -291,7 +291,7 @@ function get_value(
             level = level_boundary.level[level_boundary_idx](t + Δt)
         else
             error(
-                "Level condition node '$feature_id' is neither a basin nor a level boundary.",
+                "Level condition node '$node_id' is neither a basin nor a level boundary.",
             )
         end
 
@@ -301,7 +301,7 @@ function get_value(
         flow_boundary_idx = findsorted(flow_boundary.node_id, node_id)
 
         if flow_boundary_idx === nothing
-            error("Flow condition node #$feature_id is not a flow boundary.")
+            error("Flow condition node $node_id is not a flow boundary.")
         end
 
         value = flow_boundary.flow_rate[flow_boundary_idx](t + Δt)
@@ -529,7 +529,7 @@ function update_tabulated_rating_curve!(integrator)::Nothing
         id = first(group).node_id
         level = [row.level for row in group]
         discharge = [row.discharge for row in group]
-        i = searchsortedfirst(node_id, id)
+        i = searchsortedfirst(node_id, NodeID(id))
         tables[i] = LinearInterpolation(discharge, level; extrapolate = true)
     end
     return nothing
diff --git a/core/src/solve.jl b/core/src/solve.jl
index e7aa3acd9..7e2b3665e 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -36,7 +36,6 @@ end
 
 struct NodeMetadata
     type::Symbol
-    data_index::Int
     allocation_network_id::Int
 end
 
@@ -374,7 +373,7 @@ end
 """
 node_id: node ID of the DiscreteControl node; these are not unique but repeated
     by the amount of conditions of this DiscreteControl node
-listen_feature_id: the ID of the node/edge being condition on
+listen_node_id: the ID of the node being condition on
 variable: the name of the variable in the condition
 greater_than: The threshold value in the condition
 condition_value: The current value of each condition
@@ -501,22 +500,22 @@ function valid_n_neighbors(node::AbstractParameterNode, graph::MetaGraph)::Bool
             n_outneighbors = length(outneighbor_labels_type(graph, id, edge_type))
 
             if n_inneighbors < bounds.in_min
-                @error "Nodes of type $node_type must have at least $(bounds.in_min) $edge_type inneighbor(s) (got $n_inneighbors for node #$id)."
+                @error "Nodes of type $node_type must have at least $(bounds.in_min) $edge_type inneighbor(s) (got $n_inneighbors for node $id)."
                 errors = true
             end
 
             if n_inneighbors > bounds.in_max
-                @error "Nodes of type $node_type can have at most $(bounds.in_max) $edge_type inneighbor(s) (got $n_inneighbors for node #$id)."
+                @error "Nodes of type $node_type can have at most $(bounds.in_max) $edge_type inneighbor(s) (got $n_inneighbors for node $id)."
                 errors = true
             end
 
             if n_outneighbors < bounds.out_min
-                @error "Nodes of type $node_type must have at least $(bounds.out_min) $edge_type outneighbor(s) (got $n_outneighbors for node #$id)."
+                @error "Nodes of type $node_type must have at least $(bounds.out_min) $edge_type outneighbor(s) (got $n_outneighbors for node $id)."
                 errors = true
             end
 
             if n_outneighbors > bounds.out_max
-                @error "Nodes of type $node_type can have at most $(bounds.out_max) $edge_type outneighbor(s) (got $n_outneighbors for node #$id)."
+                @error "Nodes of type $node_type can have at most $(bounds.out_max) $edge_type outneighbor(s) (got $n_outneighbors for node $id)."
                 errors = true
             end
         end
@@ -770,8 +769,8 @@ function formulate_flow!(
     flow = get_tmp(flow, storage)
 
     for (i, id) in enumerate(node_id)
-        src_id = only(inneighbor_labels_type(graph_flow, id, EdgeType.flow))
-        dst_id = only(outneighbor_labels_type(graph_flow, id, EdgeType.flow))
+        src_id = only(inneighbor_labels_type(graph, id, EdgeType.flow))
+        dst_id = only(outneighbor_labels_type(graph, id, EdgeType.flow))
 
         if !active[i]
             continue
diff --git a/core/src/utils.jl b/core/src/utils.jl
index edd46d6d6..ea084c7cc 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -22,11 +22,11 @@ function create_graph(db::DB)::MetaGraph
         edge_data_type = EdgeMetadata,
         graph_data = nothing, # In theory all remaining connectivity data could be put here
     )
-    for (i, row) in enumerate(node_rows)
+    for row in node_rows
         allocation_network_id =
             hasproperty(row, :allocation_network_id) ? row.allocation_network_id : 0
         graph[NodeID(row.fid)] =
-            NodeMetadata(Symbol(snake_case(row.type)), i, allocation_network_id)
+            NodeMetadata(Symbol(snake_case(row.type)), allocation_network_id)
     end
     for (; from_node_id, to_node_id, edge_type, fid) in edge_rows
         if edge_type == "flow"
@@ -125,8 +125,8 @@ function get_storage_from_level(basin::Basin, state_idx::Int, level::Float64)::F
     bottom = first(level_discrete)
 
     if level < bottom
-        node_id = basin.node_id[state_idx]
-        @error "The level $level of basin #$node_id is lower than the bottom of this basin $bottom."
+        node_id = basin.node_id[NodeID(state_idx)]
+        @error "The level $level of basin $node_id is lower than the bottom of this basin $bottom."
         return NaN
     end
 
@@ -567,7 +567,7 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
 
         if !isempty(truth_states_wrong_length)
             errors = true
-            @error "DiscreteControl node #$id has $n_conditions condition(s), which is inconsistent with these truth state(s): $truth_states_wrong_length."
+            @error "DiscreteControl node $id has $n_conditions condition(s), which is inconsistent with these truth state(s): $truth_states_wrong_length."
         end
 
         # Check whether these control states are defined for the
@@ -594,7 +594,7 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
             if !isempty(undefined_control_states)
                 undefined_list = collect(undefined_control_states)
                 node_type = typeof(node).name.name
-                @error "These control states from DiscreteControl node #$id are not defined for controlled $node_type #$id_outneighbor: $undefined_list."
+                @error "These control states from DiscreteControl node $id are not defined for controlled $node_type $id_outneighbor: $undefined_list."
                 errors = true
             end
         end
@@ -606,22 +606,22 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
             # (e.g. for some node some variables are transient, some not).
             if node_type ∉ [:flow_boundary, :level_boundary]
                 errors = true
-                @error "Look ahead supplied for non-timeseries listen variable '$var' from listen node #$feature_id."
+                @error "Look ahead supplied for non-timeseries listen variable '$var' from listen node $node_id."
             else
                 if Δt < 0
                     errors = true
-                    @error "Negative look ahead supplied for listen variable '$var' from listen node #$feature_id."
+                    @error "Negative look ahead supplied for listen variable '$var' from listen node $node_id."
                 else
                     node = getfield(p, node_type)
                     idx = if node_type == :Basin
-                        id_index(node.node_id, feature_id)
+                        id_index(node.node_id, node_id)
                     else
                         searchsortedfirst(node.node_id, node_id)
                     end
                     interpolation = getfield(node, Symbol(var))[idx]
                     if t_end + Δt > interpolation.t[end]
                         errors = true
-                        @error "Look ahead for listen variable '$var' from listen node #$feature_id goes past timeseries end during simulation."
+                        @error "Look ahead for listen variable '$var' from listen node $node_id goes past timeseries end during simulation."
                     end
                 end
             end
@@ -669,7 +669,7 @@ function expand_logic_mapping(
 
                 if haskey(logic_mapping_expanded, new_key)
                     control_state_existing = logic_mapping_expanded[new_key]
-                    msg = "Multiple control states found for DiscreteControl node #$node_id for truth state `$truth_state_new`: $control_state, $control_state_existing."
+                    msg = "Multiple control states found for DiscreteControl node $node_id for truth state `$truth_state_new`: $control_state, $control_state_existing."
                     @assert control_state_existing == control_state msg
                 else
                     logic_mapping_expanded[new_key] = control_state
diff --git a/core/src/validation.jl b/core/src/validation.jl
index b3a53cf6a..29bf65dc0 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -59,37 +59,37 @@ end
 
 # Allowed types for downstream (to_node_id) nodes given the type of the upstream (from_node_id) node
 neighbortypes(nodetype::Symbol) = neighbortypes(Val(nodetype))
-neighbortypes(::Val{:Pump}) = Set((:Basin, :FractionalFlow, :Terminal, :LevelBoundary))
-neighbortypes(::Val{:Outlet}) = Set((:Basin, :FractionalFlow, :Terminal, :LevelBoundary))
-neighbortypes(::Val{:User}) = Set((:Basin, :FractionalFlow, :Terminal, :LevelBoundary))
-neighbortypes(::Val{:Basin}) = Set((
-    :LinearResistance,
-    :TabulatedRatingCurve,
-    :ManningResistance,
-    :Pump,
-    :Outlet,
-    :User,
+neighbortypes(::Val{:pump}) = Set((:basin, :fractional_flow, :terminal, :level_boundary))
+neighbortypes(::Val{:outlet}) = Set((:basin, :fractional_flow, :terminal, :level_boundary))
+neighbortypes(::Val{:user}) = Set((:basin, :fractional_flow, :terminal, :level_boundary))
+neighbortypes(::Val{:basin}) = Set((
+    :linear_resistance,
+    :tabulated_rating_curve,
+    :manning_resistance,
+    :pump,
+    :outlet,
+    :user,
 ))
-neighbortypes(::Val{:Terminal}) = Set{Symbol}() # only endnode
-neighbortypes(::Val{:FractionalFlow}) = Set((:Basin, :Terminal, :LevelBoundary))
-neighbortypes(::Val{:FlowBoundary}) =
-    Set((:Basin, :FractionalFlow, :Terminal, :LevelBoundary))
-neighbortypes(::Val{:LevelBoundary}) =
-    Set((:LinearResistance, :ManningResistance, :Pump, :Outlet))
-neighbortypes(::Val{:LinearResistance}) = Set((:Basin, :LevelBoundary))
-neighbortypes(::Val{:ManningResistance}) = Set((:Basin, :LevelBoundary))
-neighbortypes(::Val{:DiscreteControl}) = Set((
-    :Pump,
-    :Outlet,
-    :TabulatedRatingCurve,
-    :LinearResistance,
-    :ManningResistance,
-    :FractionalFlow,
-    :PidControl,
+neighbortypes(::Val{:terminal}) = Set{Symbol}() # only endnode
+neighbortypes(::Val{:fractional_flow}) = Set((:basin, :terminal, :level_boundary))
+neighbortypes(::Val{:flow_boundary}) =
+    Set((:basin, :fractional_flow, :terminal, :level_boundary))
+neighbortypes(::Val{:level_boundary}) =
+    Set((:linear_resistance, :manning_resistance, :pump, :outlet))
+neighbortypes(::Val{:linear_resistance}) = Set((:basin, :level_boundary))
+neighbortypes(::Val{:manning_resistance}) = Set((:basin, :level_boundary))
+neighbortypes(::Val{:discrete_control}) = Set((
+    :pump,
+    :outlet,
+    :tabulated_rating_curve,
+    :linear_resistance,
+    :manning_resistance,
+    :fractioal_flow,
+    :pid_control,
 ))
-neighbortypes(::Val{:PidControl}) = Set((:Pump, :Outlet))
-neighbortypes(::Val{:TabulatedRatingCurve}) =
-    Set((:Basin, :FractionalFlow, :Terminal, :LevelBoundary))
+neighbortypes(::Val{:pid_control}) = Set((:pump, :outlet))
+neighbortypes(::Val{:tabulated_rating_curve}) =
+    Set((:basin, :fractional_flow, :terminal, :level_boundary))
 neighbortypes(::Any) = Set{Symbol}()
 
 # Allowed number of inneighbors and outneighbors per node type
@@ -446,9 +446,11 @@ function valid_edges(graph::MetaGraph)::Bool
         id_dst = label_for(graph, e.dst)
         type_src = graph[id_src].type
         type_dst = graph[id_dst].type
+
         if !(type_dst in neighbortypes(type_src))
+            errors = true
             edge_id = graph[id_src, id_dst].id.value
-            @error "Cannot connect a $type_src to a $type_dst (edge #$edge_id from node #$id_src to #$id_dst)."
+            @error "Cannot connect a $type_src to a $type_dst (edge #$edge_id from node $id_src to $id_dst)."
         end
     end
     return !errors
@@ -466,20 +468,20 @@ function valid_profiles(
 
     for (id, levels, areas) in zip(node_id, level, area)
         if !allunique(levels)
-            push!(errors, "Basin #$id has repeated levels, this cannot be interpolated.")
+            push!(errors, "Basin $id has repeated levels, this cannot be interpolated.")
         end
 
         if areas[1] <= 0
             push!(
                 errors,
-                "Basin profiles cannot start with area <= 0 at the bottom for numerical reasons (got area $(areas[1]) for node #$id).",
+                "Basin profiles cannot start with area <= 0 at the bottom for numerical reasons (got area $(areas[1]) for node $id).",
             )
         end
 
         if areas[end] < areas[end - 1]
             push!(
                 errors,
-                "Basin profiles cannot have decreasing area at the top since extrapolating could lead to negative areas, found decreasing top areas for node #$id.",
+                "Basin profiles cannot have decreasing area at the top since extrapolating could lead to negative areas, found decreasing top areas for node $id.",
             )
         end
     end
@@ -509,7 +511,7 @@ function valid_flow_rates(
         if flow_rate_ < 0.0
             errors = true
             control_state = key[2]
-            @error "$node_type flow rates must be non-negative, found $flow_rate_ for control state '$control_state' of #$id_controlled."
+            @error "$node_type flow rates must be non-negative, found $flow_rate_ for control state '$control_state' of $id_controlled."
         end
     end
 
@@ -519,7 +521,7 @@ function valid_flow_rates(
         end
         if flow_rate_ < 0.0
             errors = true
-            @error "$node_type flow rates must be non-negative, found $flow_rate_ for static #$id."
+            @error "$node_type flow rates must be non-negative, found $flow_rate_ for static $id."
         end
     end
 
@@ -538,7 +540,7 @@ function valid_pid_connectivity(
     for (id, listen_id) in zip(pid_control_node_id, pid_control_listen_node_id)
         has_index, _ = id_index(basin_node_id, listen_id)
         if !has_index
-            @error "Listen node #$listen_id of PidControl node #$id is not a Basin"
+            @error "Listen node $listen_id of PidControl node $id is not a Basin"
             errors = true
         end
 
@@ -548,14 +550,14 @@ function valid_pid_connectivity(
             pump_intake_id =
                 only(inneighbor_labels_type(graph, controlled_id, EdgeType.flow))
             if pump_intake_id != listen_id
-                @error "Listen node #$listen_id of PidControl node #$id is not upstream of controlled pump #$controlled_id"
+                @error "Listen node $listen_id of PidControl node $id is not upstream of controlled pump $controlled_id"
                 errors = true
             end
         else
             outlet_outflow_id =
                 only(outneighbor_labels_type(graph, controlled_id, EdgeType.flow))
             if outlet_outflow_id != listen_id
-                @error "Listen node #$listen_id of PidControl node #$id is not downstream of controlled outlet #$controlled_id"
+                @error "Listen node $listen_id of PidControl node $id is not downstream of controlled outlet $controlled_id"
                 errors = true
             end
         end
@@ -590,7 +592,7 @@ function valid_fractional_flow(
         if src_outneighbor_ids ⊈ node_id
             errors = true
             @error(
-                "Node #$src_id combines fractional flow outneighbors with other outneigbor types."
+                "Node $src_id combines fractional flow outneighbors with other outneigbor types."
             )
         end
 
diff --git a/core/test/control.jl b/core/test/control.jl
index 7418cd142..f9214b02e 100644
--- a/core/test/control.jl
+++ b/core/test/control.jl
@@ -14,16 +14,16 @@ using PreallocationTools: get_tmp
 
     # Control input
     pump_control_mapping = p.pump.control_mapping
-    @test pump_control_mapping[(4, "off")].flow_rate == 0
-    @test pump_control_mapping[(4, "on")].flow_rate == 1.0e-5
-
-    logic_mapping::Dict{Tuple{Int, String}, String} = Dict(
-        (5, "TT") => "on",
-        (6, "F") => "active",
-        (5, "TF") => "off",
-        (5, "FF") => "on",
-        (5, "FT") => "off",
-        (6, "T") => "inactive",
+    @test pump_control_mapping[(Ribasim.NodeID(4), "off")].flow_rate == 0
+    @test pump_control_mapping[(Ribasim.NodeID(4), "on")].flow_rate == 1.0e-5
+
+    logic_mapping::Dict{Tuple{Ribasim.NodeID, String}, String} = Dict(
+        (Ribasim.NodeID(5), "TT") => "on",
+        (Ribasim.NodeID(6), "F") => "active",
+        (Ribasim.NodeID(5), "TF") => "off",
+        (Ribasim.NodeID(5), "FF") => "on",
+        (Ribasim.NodeID(5), "FT") => "off",
+        (Ribasim.NodeID(6), "T") => "inactive",
     )
 
     @test discrete_control.logic_mapping == logic_mapping
@@ -196,8 +196,9 @@ end
     timesteps = Ribasim.timesteps(model)
     level = Ribasim.get_storages_and_levels(model).level[1, :]
 
-    target_high = pid_control.control_mapping[(6, "target_high")].target.u[1]
-    target_low = pid_control.control_mapping[(6, "target_low")].target.u[1]
+    target_high =
+        pid_control.control_mapping[(Ribasim.NodeID(6), "target_high")].target.u[1]
+    target_low = pid_control.control_mapping[(Ribasim.NodeID(6), "target_low")].target.u[1]
 
     t_target_jump = discrete_control.record.time[2]
     t_idx_target_jump = searchsortedlast(timesteps, t_target_jump)
diff --git a/core/test/run_models.jl b/core/test/run_models.jl
index 293a623f1..9a73267a6 100644
--- a/core/test/run_models.jl
+++ b/core/test/run_models.jl
@@ -60,8 +60,8 @@ end
     t = table.time[1]
     @test length(p.fractional_flow.node_id) == 3
     for id in p.fractional_flow.node_id
-        inflow = only(table.flow[table.to_node_id .== id .&& table.time .== t])
-        outflow = only(table.flow[table.from_node_id .== id .&& table.time .== t])
+        inflow = only(table.flow[table.to_node_id .== id.value .&& table.time .== t])
+        outflow = only(table.flow[table.from_node_id .== id.value .&& table.time .== t])
         @test inflow == outflow
     end
 end
diff --git a/core/test/utils.jl b/core/test/utils.jl
index 1237d4f92..1f7e581d4 100644
--- a/core/test/utils.jl
+++ b/core/test/utils.jl
@@ -7,9 +7,9 @@ using SQLite
 using Logging
 
 @testset "id_index" begin
-    ids = Indices([2, 4, 6])
-    @test Ribasim.id_index(ids, 4) === (true, 2)
-    @test Ribasim.id_index(ids, 5) === (false, 0)
+    ids = Indices([Ribasim.NodeID(2), Ribasim.NodeID(4), Ribasim.NodeID(6)])
+    @test Ribasim.id_index(ids, Ribasim.NodeID(4)) === (true, 2)
+    @test Ribasim.id_index(ids, Ribasim.NodeID(5)) === (false, 0)
 end
 
 @testset "profile_storage" begin
@@ -26,7 +26,7 @@ end
     darea = zeros(2)
     storage = Ribasim.profile_storage.(level, area)
     basin = Ribasim.Basin(
-        Indices([5, 7]),
+        Indices([Ribasim.NodeID(5), Ribasim.NodeID(7)]),
         [2.0, 3.0],
         [2.0, 3.0],
         [2.0, 3.0],
@@ -40,17 +40,32 @@ end
     )
 
     @test basin.level[2][1] === 4.0
-    @test Ribasim.basin_bottom(basin, 5) === 0.0
-    @test Ribasim.basin_bottom(basin, 7) === 4.0
-    @test Ribasim.basin_bottom(basin, 6) === nothing
-    @test Ribasim.basin_bottoms(basin, 5, 7, 6) === (0.0, 4.0)
-    @test Ribasim.basin_bottoms(basin, 5, 0, 6) === (0.0, 0.0)
-    @test Ribasim.basin_bottoms(basin, 0, 7, 6) === (4.0, 4.0)
-    @test_throws "No bottom defined on either side of 6" Ribasim.basin_bottoms(
+    @test Ribasim.basin_bottom(basin, Ribasim.NodeID(5)) === 0.0
+    @test Ribasim.basin_bottom(basin, Ribasim.NodeID(7)) === 4.0
+    @test Ribasim.basin_bottom(basin, Ribasim.NodeID(6)) === nothing
+    @test Ribasim.basin_bottoms(
         basin,
-        0,
-        1,
-        6,
+        Ribasim.NodeID(5),
+        Ribasim.NodeID(7),
+        Ribasim.NodeID(6),
+    ) === (0.0, 4.0)
+    @test Ribasim.basin_bottoms(
+        basin,
+        Ribasim.NodeID(5),
+        Ribasim.NodeID(0),
+        Ribasim.NodeID(6),
+    ) === (0.0, 0.0)
+    @test Ribasim.basin_bottoms(
+        basin,
+        Ribasim.NodeID(0),
+        Ribasim.NodeID(7),
+        Ribasim.NodeID(6),
+    ) === (4.0, 4.0)
+    @test_throws "No bottom defined on either side of #6" Ribasim.basin_bottoms(
+        basin,
+        Ribasim.NodeID(0),
+        Ribasim.NodeID(1),
+        Ribasim.NodeID(6),
     )
 end
 
@@ -81,7 +96,7 @@ end
     ]
     storage = Ribasim.profile_storage(level, area)
     basin = Ribasim.Basin(
-        Indices([1]),
+        Indices([Ribasim.NodeID(1)]),
         zeros(1),
         zeros(1),
         zeros(1),
@@ -115,19 +130,19 @@ end
 end
 
 @testset "Expand logic_mapping" begin
-    logic_mapping = Dict{Tuple{Int, String}, String}()
-    logic_mapping[(1, "*T*")] = "foo"
-    logic_mapping[(2, "FF")] = "bar"
+    logic_mapping = Dict{Tuple{Ribasim.NodeID, String}, String}()
+    logic_mapping[(Ribasim.NodeID(1), "*T*")] = "foo"
+    logic_mapping[(Ribasim.NodeID(2), "FF")] = "bar"
     logic_mapping_expanded = Ribasim.expand_logic_mapping(logic_mapping)
 
-    @test logic_mapping_expanded[(1, "TTT")] == "foo"
-    @test logic_mapping_expanded[(1, "FTT")] == "foo"
-    @test logic_mapping_expanded[(1, "TTF")] == "foo"
-    @test logic_mapping_expanded[(1, "FTF")] == "foo"
-    @test logic_mapping_expanded[(2, "FF")] == "bar"
+    @test logic_mapping_expanded[(Ribasim.NodeID(1), "TTT")] == "foo"
+    @test logic_mapping_expanded[(Ribasim.NodeID(1), "FTT")] == "foo"
+    @test logic_mapping_expanded[(Ribasim.NodeID(1), "TTF")] == "foo"
+    @test logic_mapping_expanded[(Ribasim.NodeID(1), "FTF")] == "foo"
+    @test logic_mapping_expanded[(Ribasim.NodeID(2), "FF")] == "bar"
     @test length(logic_mapping_expanded) == 5
 
-    new_key = (3, "duck")
+    new_key = (Ribasim.NodeID(3), "duck")
     logic_mapping[new_key] = "quack"
 
     @test_throws "Truth state 'duck' contains illegal characters or is empty." Ribasim.expand_logic_mapping(
@@ -136,7 +151,7 @@ end
 
     delete!(logic_mapping, new_key)
 
-    new_key = (3, "")
+    new_key = (Ribasim.NodeID(3), "")
     logic_mapping[new_key] = "bar"
 
     @test_throws "Truth state '' contains illegal characters or is empty." Ribasim.expand_logic_mapping(
@@ -145,13 +160,13 @@ end
 
     delete!(logic_mapping, new_key)
 
-    new_key = (1, "FTT")
+    new_key = (Ribasim.NodeID(1), "FTT")
     logic_mapping[new_key] = "foo"
 
     # This should not throw an error, as although "FTT" for node_id = 1 is already covered above, this is consistent
     Ribasim.expand_logic_mapping(logic_mapping)
 
-    new_key = (1, "TTF")
+    new_key = (Ribasim.NodeID(1), "TTF")
     logic_mapping[new_key] = "bar"
 
     @test_throws "Multiple control states found for DiscreteControl node #1 for truth state `TTF`: foo, bar." Ribasim.expand_logic_mapping(
diff --git a/core/test/validation.jl b/core/test/validation.jl
index a193cb816..616a57969 100644
--- a/core/test/validation.jl
+++ b/core/test/validation.jl
@@ -6,9 +6,10 @@ using DataInterpolations: LinearInterpolation
 import SQLite
 using Logging
 using Test
+using MetaGraphsNext: MetaGraph
 
 @testset "Basin profile validation" begin
-    node_id = Indices([1])
+    node_id = Indices([Ribasim.NodeID(1)])
     level = [[0.0, 0.0, 1.0]]
     area = [[0.0, 100.0, 90]]
     errors = Ribasim.valid_profiles(node_id, level, area)
@@ -52,27 +53,41 @@ end
 end
 
 @testset "Neighbor count validation" begin
-    graph_flow = DiGraph(6)
-    add_edge!(graph_flow, 2, 1)
-    add_edge!(graph_flow, 3, 1)
-    add_edge!(graph_flow, 6, 2)
+    graph = MetaGraph(
+        DiGraph();
+        label_type = Ribasim.NodeID,
+        vertex_data_type = Ribasim.NodeMetadata,
+        edge_data_type = Ribasim.EdgeMetadata,
+        graph_data = nothing,
+    )
+
+    for i in 1:6
+        type = i in [1, 6] ? :pump : :other
+        graph[Ribasim.NodeID(i)] = Ribasim.NodeMetadata(type, 9)
+    end
 
-    graph_control = DiGraph(6)
-    add_edge!(graph_control, 5, 6)
+    graph[Ribasim.NodeID(2), Ribasim.NodeID(1)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(1), Ribasim.EdgeType.flow)
+    graph[Ribasim.NodeID(3), Ribasim.NodeID(1)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(2), Ribasim.EdgeType.flow)
+    graph[Ribasim.NodeID(6), Ribasim.NodeID(2)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(3), Ribasim.EdgeType.flow)
+    graph[Ribasim.NodeID(5), Ribasim.NodeID(6)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(4), Ribasim.EdgeType.control)
 
     pump = Ribasim.Pump(
-        [1, 6],
+        [Ribasim.NodeID(1), Ribasim.NodeID(6)],
         [true, true],
         [0.0, 0.0],
         [0.0, 0.0],
         [1.0, 1.0],
-        Dict{Tuple{Int, String}, NamedTuple}(),
+        Dict{Tuple{Ribasim.NodeID, String}, NamedTuple}(),
         falses(2),
     )
 
     logger = TestLogger()
     with_logger(logger) do
-        @test !Ribasim.valid_n_neighbors(pump, graph_flow, graph_control)
+        @test !Ribasim.valid_n_neighbors(pump, graph)
     end
 
     @test length(logger.logs) == 3
@@ -86,16 +101,22 @@ end
     @test logger.logs[3].message ==
           "Nodes of type Ribasim.Pump{Vector{Float64}} must have at least 1 flow inneighbor(s) (got 0 for node #6)."
 
-    add_edge!(graph_flow, 2, 5)
-    add_edge!(graph_flow, 5, 3)
-    add_edge!(graph_flow, 5, 4)
+    graph[Ribasim.NodeID(2), Ribasim.NodeID(5)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(5), Ribasim.EdgeType.flow)
+    graph[Ribasim.NodeID(5), Ribasim.NodeID(3)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(6), Ribasim.EdgeType.flow)
+    graph[Ribasim.NodeID(5), Ribasim.NodeID(4)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(7), Ribasim.EdgeType.flow)
 
-    fractional_flow =
-        Ribasim.FractionalFlow([5], [1.0], Dict{Tuple{Int, String}, NamedTuple}())
+    fractional_flow = Ribasim.FractionalFlow(
+        [Ribasim.NodeID(5)],
+        [1.0],
+        Dict{Tuple{Int, String}, NamedTuple}(),
+    )
 
     logger = TestLogger()
     with_logger(logger) do
-        @test !Ribasim.valid_n_neighbors(fractional_flow, graph_flow, graph_control)
+        @test !Ribasim.valid_n_neighbors(fractional_flow, graph)
     end
 
     @test length(logger.logs) == 2
@@ -115,28 +136,44 @@ end
 end
 
 @testset "PidControl connectivity validation" begin
-    pid_control_node_id = [1, 6]
-    pid_control_listen_node_id = [3, 5]
-    pump_node_id = [2, 4]
+    pid_control_node_id = Ribasim.NodeID.([1, 6])
+    pid_control_listen_node_id = Ribasim.NodeID.([3, 5])
+    pump_node_id = Ribasim.NodeID.([2, 4])
+
+    graph = MetaGraph(
+        DiGraph();
+        label_type = Ribasim.NodeID,
+        vertex_data_type = Ribasim.NodeMetadata,
+        edge_data_type = Ribasim.EdgeMetadata,
+        graph_data = nothing,
+    )
 
-    graph_flow = DiGraph(7)
-    graph_control = DiGraph(7)
+    graph[Ribasim.NodeID(1)] = Ribasim.NodeMetadata(:pid_control, 0)
+    graph[Ribasim.NodeID(6)] = Ribasim.NodeMetadata(:pid_control, 0)
+    graph[Ribasim.NodeID(2)] = Ribasim.NodeMetadata(:pump, 0)
+    graph[Ribasim.NodeID(4)] = Ribasim.NodeMetadata(:pump, 0)
+    graph[Ribasim.NodeID(3)] = Ribasim.NodeMetadata(:something_else, 0)
+    graph[Ribasim.NodeID(5)] = Ribasim.NodeMetadata(:basin, 0)
+    graph[Ribasim.NodeID(7)] = Ribasim.NodeMetadata(:basin, 0)
 
-    add_edge!(graph_flow, 3, 4)
-    add_edge!(graph_flow, 7, 2)
+    graph[Ribasim.NodeID(3), Ribasim.NodeID(4)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(1), Ribasim.EdgeType.flow)
+    graph[Ribasim.NodeID(7), Ribasim.NodeID(2)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(2), Ribasim.EdgeType.flow)
 
-    add_edge!(graph_control, 1, 4)
-    add_edge!(graph_control, 6, 2)
+    graph[Ribasim.NodeID(1), Ribasim.NodeID(4)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(3), Ribasim.EdgeType.control)
+    graph[Ribasim.NodeID(6), Ribasim.NodeID(2)] =
+        Ribasim.EdgeMetadata(Ribasim.EdgeID(4), Ribasim.EdgeType.control)
 
-    basin_node_id = Indices([5, 7])
+    basin_node_id = Indices(Ribasim.NodeID.([5, 7]))
 
     logger = TestLogger()
     with_logger(logger) do
         @test !Ribasim.valid_pid_connectivity(
             pid_control_node_id,
             pid_control_listen_node_id,
-            graph_flow,
-            graph_control,
+            graph,
             basin_node_id,
             pump_node_id,
         )
@@ -168,7 +205,7 @@ if !Sys.islinux()
         logger = TestLogger()
         with_logger(logger) do
             @test !Ribasim.valid_fractional_flow(
-                connectivity.graph_flow,
+                connectivity.graph,
                 fractional_flow.node_id,
                 fractional_flow.control_mapping,
             )
@@ -181,13 +218,13 @@ if !Sys.islinux()
         @test logger.logs[2].level == Error
         @test logger.logs[2].message ==
               "Fractional flow nodes must have non-negative fractions."
-        @test logger.logs[2].kwargs[:node_id] == 3
+        @test logger.logs[2].kwargs[:node_id].value == 3
         @test logger.logs[2].kwargs[:fraction] ≈ -0.1
         @test logger.logs[2].kwargs[:control_state] == ""
         @test logger.logs[3].level == Error
         @test logger.logs[3].message ==
               "The sum of fractional flow fractions leaving a node must be ≈1."
-        @test logger.logs[3].kwargs[:node_id] == 7
+        @test logger.logs[3].kwargs[:node_id].value == 7
         @test logger.logs[3].kwargs[:fraction_sum] ≈ 0.4
         @test logger.logs[3].kwargs[:control_state] == ""
     end
@@ -233,13 +270,13 @@ end
 
     with_logger(logger) do
         @test_throws "Invalid Outlet flow rate(s)." Ribasim.Outlet(
-            [1],
+            [Ribasim.NodeID(1)],
             [true],
             [-1.0],
             [NaN],
             [NaN],
             [NaN],
-            Dict{Tuple{Int, String}, NamedTuple}(),
+            Dict{Tuple{Ribasim.NodeID, String}, NamedTuple}(),
             [false],
         )
     end
@@ -253,12 +290,14 @@ end
 
     with_logger(logger) do
         @test_throws "Invalid Pump flow rate(s)." Ribasim.Pump(
-            [1],
+            [Ribasim.NodeID(1)],
             [true],
             [-1.0],
             [NaN],
             [NaN],
-            Dict{Tuple{Int, String}, NamedTuple}((1, "foo") => (; flow_rate = -1.0)),
+            Dict{Tuple{Ribasim.NodeID, String}, NamedTuple}(
+                (Ribasim.NodeID(1), "foo") => (; flow_rate = -1.0),
+            ),
             [false],
         )
     end

From 68c66ae6d82a768d31006e2b6cd814f36dc5a3b7 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Fri, 17 Nov 2023 14:53:52 +0100
Subject: [PATCH 6/8] Docstrings

---
 core/src/solve.jl |  6 ++++++
 core/src/utils.jl | 34 +++++++++++++++++++++++++++++-----
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/core/src/solve.jl b/core/src/solve.jl
index 7e2b3665e..23614a073 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -34,11 +34,17 @@ end
 
 @enumx EdgeType flow control
 
+"""
+Type for storing metadata of nodes in the graph.
+"""
 struct NodeMetadata
     type::Symbol
     allocation_network_id::Int
 end
 
+"""
+Type for storing metadata of edges in the graph.
+"""
 struct EdgeMetadata
     id::EdgeID
     type::EdgeType.T
diff --git a/core/src/utils.jl b/core/src/utils.jl
index ea084c7cc..fa0b1446b 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -12,6 +12,15 @@ function valid_edge_types(db::DB)::Bool
     return !errors
 end
 
+"""
+Return a directed metagraph with data of nodes (NodeMetadata):
+- Node type (snake case)
+- Allocation network ID
+
+and data of edges (EdgeMetadata):
+- Edge ID (EdgeID)
+- type (flow/control)
+"""
 function create_graph(db::DB)::MetaGraph
     node_rows = execute(db, "select fid, type from Node")
     edge_rows = execute(db, "select fid, from_node_id, to_node_id, edge_type from Edge")
@@ -29,11 +38,10 @@ function create_graph(db::DB)::MetaGraph
             NodeMetadata(Symbol(snake_case(row.type)), allocation_network_id)
     end
     for (; from_node_id, to_node_id, edge_type, fid) in edge_rows
-        if edge_type == "flow"
-            edge_type = EdgeType.flow
-        elseif edge_type == "control"
-            edge_type = EdgeType.control
-        else
+        try
+            # hasfield does not work
+            edge_type = getfield(EdgeType, Symbol(edge_type))
+        catch
             error("Invalid edge type $edge_type.")
         end
         graph[NodeID(from_node_id), NodeID(to_node_id)] =
@@ -42,6 +50,10 @@ function create_graph(db::DB)::MetaGraph
     return graph
 end
 
+"""
+Get the inneighbor node IDs of the given node ID (label)
+over the given edge type in the graph.
+"""
 function inneighbor_labels_type(
     graph::MetaGraph,
     label::NodeID,
@@ -53,6 +65,10 @@ function inneighbor_labels_type(
     ]
 end
 
+"""
+Get the outneighbor node IDs of the given node ID (label)
+over the given edge type in the graph.
+"""
 function outneighbor_labels_type(
     graph::MetaGraph,
     label::NodeID,
@@ -64,6 +80,10 @@ function outneighbor_labels_type(
     ]
 end
 
+"""
+Get the in- and outneighbor node IDs of the given node ID (label)
+over the given edge type in the graph.
+"""
 function all_neighbor_labels_type(
     graph::MetaGraph,
     label::NodeID,
@@ -75,6 +95,10 @@ function all_neighbor_labels_type(
     ]
 end
 
+"""
+Get the metadata of an edge in the graph from an edge of the underlying
+DiGraph.
+"""
 function metadata_from_edge(graph::MetaGraph, edge::Edge{Int})::EdgeMetadata
     label_src = label_for(graph, edge.src)
     label_dst = label_for(graph, edge.dst)

From 4481046dcfb411417f8cd85f494b25b062008500 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Fri, 17 Nov 2023 14:57:42 +0100
Subject: [PATCH 7/8] Update docs

---
 docs/contribute/addnode.qmd | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/docs/contribute/addnode.qmd b/docs/contribute/addnode.qmd
index 755aea615..df21af083 100644
--- a/docs/contribute/addnode.qmd
+++ b/docs/contribute/addnode.qmd
@@ -12,7 +12,7 @@ The parameters object (defined in `solve.jl`) passed to the ODE solver must be m
 
 ```julia
 struct NewNodeType
-    node_id::Vector{Int}
+    node_id::Vector{NodeID}
     # Other fields
 end
 ```
@@ -57,7 +57,7 @@ function NewNodeType(db::DB, config::Config)::NewNodeType
 
     # Unpack the fields of static as inputs for the NewNodeType constructor
     return NewNodeType(
-        parsed_parameters.node_id,
+        NodeID.(parsed_parameters.node_id),
         parsed_parameters.some_property,
         parsed_parameters.control_mapping)
 end
@@ -189,13 +189,13 @@ class NewNodeTypeStatic:
 
 The new node type might have associated restrictions for a model with the new node type so that it behaves properly. Basic node ID and node type validation happens in `Model.validate_model` in `python/ribasim/ribasim/model.py`, which automatically considers all node types in the `node_types` module.
 
-Connectivity validation happens in `valid_edges` and `valid_n_flow_neighbors` in `core/src/solve.jl`. Connectivity rules are specified in `core/src/validation.jl`. Allowed upstream and downstream neighbor types for `NewNodeType` in are specified as follows:
+Connectivity validation happens in `valid_edges` and `valid_n_flow_neighbors` in `core/src/solve.jl`. Connectivity rules are specified in `core/src/validation.jl`. Allowed upstream and downstream neighbor types for `new_node_type` (the snake case version of `NewNodeType`) are specified as follows:
 
 ```julia
 # set allowed downstream types
-neighbortypes(::Val{:NewNodeType}) = Set((:Basin,))
+neighbortypes(::Val{:new_node_type}) = Set((:basin,))
 # add your newnodetype as acceptable downstream connection of other types
-neighbortypes(::Val{:Pump}) = Set((:Basin, :NewNodeType))
+neighbortypes(::Val{:pump}) = Set((:basin, :new_node_type))
 ```
 
 The minimum and maximum allowed number of inneighbors and outneighbors for `NewNodeType` are specified as follows:

From 33a3fb786ea10df4e609ed28d5f0836b95af7aed Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 20 Nov 2023 14:58:46 +0100
Subject: [PATCH 8/8] Fix tests

---
 core/src/solve.jl            | 10 +++++++---
 core/src/validation.jl       |  6 +++---
 core/test/validation_test.jl |  8 +++++---
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/core/src/solve.jl b/core/src/solve.jl
index 23614a073..845b67754 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -53,7 +53,12 @@ end
 """
 Store the connectivity information
 
-graph_flow, graph_control: directed graph with vertices equal to ids
+graph: a directed metagraph with data of nodes (NodeMetadata):
+  - Node type (snake case)
+  - Allocation network ID
+  and data of edges (EdgeMetadata):
+  - Edge ID (EdgeID)
+  - type (flow/control)
 flow: store the flow on every flow edge
 edge_ids_flow, edge_ids_control: get the external edge id from (src, dst)
 edge_connection_type_flow, edge_connection_types_control: get (src_node_type, dst_node_type) from edge id
@@ -744,8 +749,7 @@ function continuous_control!(
         if controls_pump
             for id in outneighbor_labels_type(graph, controlled_node_id, EdgeType.flow)
                 if id in fractional_flow.node_id
-                    after_ff_id =
-                        only(outneighbor_labels_type(graph_flow, id, EdgeType.flow))
+                    after_ff_id = only(outneighbor_labels_type(graph, id, EdgeType.flow))
                     ff_idx = findsorted(fractional_flow, id)
                     flow_rate_fraction = fractional_flow.fraction[ff_idx] * flow_rate
                     flow[id, after_ff_id] = flow_rate_fraction
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 29bf65dc0..d58025454 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -571,7 +571,7 @@ Check that nodes that have fractional flow outneighbors do not have any other ty
 outneighbor, that the fractions leaving a node add up to ≈1 and that the fractions are non-negative.
 """
 function valid_fractional_flow(
-    graph_flow::MetaGraph,
+    graph::MetaGraph,
     node_id::Vector{NodeID},
     control_mapping::Dict{Tuple{NodeID, String}, NamedTuple},
 )::Bool
@@ -581,14 +581,14 @@ function valid_fractional_flow(
     src_ids = Set{NodeID}()
 
     for id in node_id
-        union!(src_ids, inneighbor_labels(graph_flow, id))
+        union!(src_ids, inneighbor_labels(graph, id))
     end
 
     node_id_set = Set{NodeID}(node_id)
     control_states = Set{String}([key[2] for key in keys(control_mapping)])
 
     for src_id in src_ids
-        src_outneighbor_ids = Set(outneighbor_labels(graph_flow, src_id))
+        src_outneighbor_ids = Set(outneighbor_labels(graph, src_id))
         if src_outneighbor_ids ⊈ node_id
             errors = true
             @error(
diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index deea526ea..9d16ac4f3 100644
--- a/core/test/validation_test.jl
+++ b/core/test/validation_test.jl
@@ -49,6 +49,7 @@ end
 end
 
 @testitem "Neighbor count validation" begin
+    using MetaGraphsNext: MetaGraph
     using Graphs: DiGraph
     using Logging
 
@@ -135,6 +136,7 @@ end
 end
 
 @testitem "PidControl connectivity validation" begin
+    using MetaGraphsNext: MetaGraph
     using Graphs: DiGraph
     using Dictionaries: Indices
     using Logging
@@ -209,7 +211,7 @@ end
     logger = TestLogger()
     with_logger(logger) do
         @test !Ribasim.valid_fractional_flow(
-            connectivity.graph_flow,
+            connectivity.graph,
             fractional_flow.node_id,
             fractional_flow.control_mapping,
         )
@@ -222,13 +224,13 @@ end
     @test logger.logs[2].level == Error
     @test logger.logs[2].message ==
           "Fractional flow nodes must have non-negative fractions."
-    @test logger.logs[2].kwargs[:node_id] == 3
+    @test logger.logs[2].kwargs[:node_id].value == 3
     @test logger.logs[2].kwargs[:fraction] ≈ -0.1
     @test logger.logs[2].kwargs[:control_state] == ""
     @test logger.logs[3].level == Error
     @test logger.logs[3].message ==
           "The sum of fractional flow fractions leaving a node must be ≈1."
-    @test logger.logs[3].kwargs[:node_id] == 7
+    @test logger.logs[3].kwargs[:node_id].value == 7
     @test logger.logs[3].kwargs[:fraction_sum] ≈ 0.4
     @test logger.logs[3].kwargs[:control_state] == ""
 end