Increase performance by reducing lookups

core/src/graph.jl

@@ -21,7 +21,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
     # The number of flow edges
     flow_counter = 0
     # Dictionary from flow edge to index in flow vector
-    flow_dict = Dict{Tuple{NodeID, NodeID}, Int}()
+    flow_dict = Dict{Tuple{NodeID, NodeID}, Int32}()
     graph = MetaGraph(
         DiGraph();
         label_type = NodeID,
@@ -65,7 +65,14 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         if ismissing(subnetwork_id)
             subnetwork_id = 0
         end
-        edge_metadata = EdgeMetadata(fid, edge_type, subnetwork_id, (id_src, id_dst))
+        edge_metadata = EdgeMetadata(
+            fid,
+            edge_type == EdgeType.flow ? flow_counter + 1 : 0,
+            edge_type,
+            subnetwork_id,
+            (id_src, id_dst),
+            (0, 0),
+        )
         if haskey(graph, id_src, id_dst)
             errors = true
             @error "Duplicate edge" id_src id_dst
@@ -169,8 +176,20 @@ end
 Set the given flow q over the edge between the given nodes.
 """
 function set_flow!(graph::MetaGraph, id_src::NodeID, id_dst::NodeID, q::Number)::Nothing
-    (; flow_dict, flow) = graph[]
-    get_tmp(flow, q)[flow_dict[(id_src, id_dst)]] = q
+    (; flow_dict) = graph[]
+    flow_idx = flow_dict[(id_src, id_dst)]
+    set_flow!(graph, flow_idx, q)
+    return nothing
+end
+
+function set_flow!(graph::MetaGraph, edge_metadata::EdgeMetadata, q::Number)::Nothing
+    set_flow!(graph, edge_metadata.flow_idx, q)
+    return nothing
+end
+
+function set_flow!(graph, flow_idx::Int32, q::Number)::Nothing
+    (; flow) = graph[]
+    get_tmp(flow, q)[flow_idx] = q
     return nothing
 end
 
@@ -184,9 +203,19 @@ function get_flow(
     val;
     prev::Bool = false,
 )::Number
-    (; flow_dict, flow, flow_prev) = graph[]
+    (; flow_dict) = graph[]
+    flow_idx = flow_dict[id_src, id_dst]
+    return get_flow(graph, flow_idx, val; prev)
+end
+
+function get_flow(graph, edge_metadata::EdgeMetadata, val; prev::Bool = false)::Number
+    return get_flow(graph, edge_metadata.flow_idx, val; prev)
+end
+
+function get_flow(graph::MetaGraph, flow_idx::Int32, val; prev::Bool = false)
+    (; flow, flow_prev) = graph[]
     flow_vector = prev ? flow_prev : flow
-    return get_tmp(flow_vector, val)[flow_dict[id_src, id_dst]]
+    return get_tmp(flow_vector, val)[flow_idx]
 end
 
 """

core/src/parameter.jl

@@ -116,16 +116,20 @@ end
 """
 Type for storing metadata of edges in the graph:
 id: ID of the edge (only used for labeling flow output)
+flow_idx: Index in the vector of flows
 type: type of the edge
 subnetwork_id_source: ID of subnetwork where this edge is a source
   (0 if not a source)
 edge: (from node ID, to node ID)
+basin_idxs: Basin indices of source and destination nodes (0 if not a basin)
 """
 struct EdgeMetadata
     id::Int32
+    flow_idx::Int32
     type::EdgeType.T
     subnetwork_id_source::Int32
     edge::Tuple{NodeID, NodeID}
+    basin_idxs::Tuple{Int32, Int32}
 end
 
 abstract type AbstractParameterNode end
@@ -233,17 +237,17 @@ Type parameter C indicates the content backing the StructVector, which can be a
 of Vectors or Arrow Primitives, and is added to avoid type instabilities.
 
 node_id: node ID of the TabulatedRatingCurve node
-inflow_id: node ID across the incoming flow edge
-outflow_ids: node IDs across the outgoing flow edges
+inflow_edge: incoming flow edge metadata
+outflow_edges: outgoing flow edges metadata
 active: whether this node is active and thus contributes flows
 tables: The current Q(h) relationships
 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{NodeID}
-    inflow_id::Vector{NodeID}
-    outflow_ids::Vector{Vector{NodeID}}
+    inflow_edge::Vector{EdgeMetadata}
+    outflow_edges::Vector{Vector{EdgeMetadata}}
     active::BitVector
     tables::Vector{ScalarInterpolation}
     time::StructVector{TabulatedRatingCurveTimeV1, C, Int}
@@ -252,17 +256,17 @@ end
 
 """
 node_id: node ID of the LinearResistance node
-inflow_id: node ID across the incoming flow edge
-outflow_id: node ID across the outgoing flow edge
+inflow_edge: incoming flow edge metadata
+outflow_edge: outgoing flow edge metadata
 active: whether this node is active and thus contributes flows
 resistance: the resistance to flow; `Q_unlimited = Δh/resistance`
 max_flow_rate: the maximum flow rate allowed through the node; `Q = clamp(Q_unlimited, -max_flow_rate, max_flow_rate)`
 control_mapping: dictionary from (node_id, control_state) to resistance and/or active state
 """
 struct LinearResistance <: AbstractParameterNode
     node_id::Vector{NodeID}
-    inflow_id::Vector{NodeID}
-    outflow_id::Vector{NodeID}
+    inflow_edge::Vector{EdgeMetadata}
+    outflow_edge::Vector{EdgeMetadata}
     active::BitVector
     resistance::Vector{Float64}
     max_flow_rate::Vector{Float64}
@@ -273,8 +277,8 @@ end
 This is a simple Manning-Gauckler reach connection.
 
 node_id: node ID of the ManningResistance node
-inflow_id: node ID across the incoming flow edge
-outflow_id: node ID across the outgoing flow edge
+inflow_edge: incoming flow edge metadata
+outflow_edge: outgoing flow edge metadata
 length: reach length
 manning_n: roughness; Manning's n in (SI units).
 
@@ -307,13 +311,15 @@ Requirements:
 """
 struct ManningResistance <: AbstractParameterNode
     node_id::Vector{NodeID}
-    inflow_id::Vector{NodeID}
-    outflow_id::Vector{NodeID}
+    inflow_edge::Vector{EdgeMetadata}
+    outflow_edge::Vector{EdgeMetadata}
     active::BitVector
     length::Vector{Float64}
     manning_n::Vector{Float64}
     profile_width::Vector{Float64}
     profile_slope::Vector{Float64}
+    upstream_bottom::Vector{Float64}
+    downstream_bottom::Vector{Float64}
     control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
 end
 
@@ -325,15 +331,15 @@ Requirements:
 * fraction must be positive.
 
 node_id: node ID of the TabulatedRatingCurve node
-inflow_id: node ID across the incoming flow edge
-outflow_id: node ID across the outgoing flow edge
+inflow_edge: incoming flow edge metadata
+outflow_edge: outgoing flow edge metadata
 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{NodeID}
-    inflow_id::Vector{NodeID}
-    outflow_id::Vector{NodeID}
+    inflow_edge::Vector{EdgeMetadata}
+    outflow_edge::Vector{EdgeMetadata}
     fraction::Vector{Float64}
     control_mapping::Dict{Tuple{NodeID, String}, NamedTuple}
 end
@@ -362,8 +368,8 @@ end
 
 """
 node_id: node ID of the Pump node
-inflow_id: node ID across the incoming flow edge
-outflow_ids: node IDs across the outgoing flow edges
+inflow_edge: incoming flow edge metadata
+outflow_edges: outgoing flow edges metadata
 active: whether this node is active and thus contributes flow
 flow_rate: target flow rate
 min_flow_rate: The minimal flow rate of the pump
@@ -373,8 +379,8 @@ is_pid_controlled: whether the flow rate of this pump is governed by PID control
 """
 struct Pump{T} <: AbstractParameterNode
     node_id::Vector{NodeID}
-    inflow_id::Vector{NodeID}
-    outflow_ids::Vector{Vector{NodeID}}
+    inflow_edge::Vector{EdgeMetadata}
+    outflow_edges::Vector{Vector{EdgeMetadata}}
     active::BitVector
     flow_rate::T
     min_flow_rate::Vector{Float64}
@@ -413,8 +419,8 @@ end
 
 """
 node_id: node ID of the Outlet node
-inflow_id: node ID across the incoming flow edge
-outflow_ids: node IDs across the outgoing flow edges
+inflow_edge: incoming flow edge metadata
+outflow_edges: outgoing flow edges metadata
 active: whether this node is active and thus contributes flow
 flow_rate: target flow rate
 min_flow_rate: The minimal flow rate of the outlet
@@ -424,8 +430,8 @@ is_pid_controlled: whether the flow rate of this outlet is governed by PID contr
 """
 struct Outlet{T} <: AbstractParameterNode
     node_id::Vector{NodeID}
-    inflow_id::Vector{NodeID}
-    outflow_ids::Vector{Vector{NodeID}}
+    inflow_edge::Vector{EdgeMetadata}
+    outflow_edges::Vector{Vector{EdgeMetadata}}
     active::BitVector
     flow_rate::T
     min_flow_rate::Vector{Float64}
@@ -528,8 +534,8 @@ end
 
 """
 node_id: node ID of the UserDemand node
-inflow_id: node ID across the incoming flow edge
-outflow_id: node ID across the outgoing flow edge
+inflow_edge: incoming flow edge metadata
+outflow_edge: outgoing flow edge metadata
 active: whether this node is active and thus demands water
 realized_bmi: Cumulative inflow volume, for read or reset by BMI only
 demand: water flux demand of UserDemand per priority over time
@@ -545,8 +551,8 @@ min_level: The level of the source basin below which the UserDemand does not abs
 """
 struct UserDemand <: AbstractParameterNode
     node_id::Vector{NodeID}
-    inflow_id::Vector{NodeID}
-    outflow_id::Vector{NodeID}
+    inflow_edge::Vector{EdgeMetadata}
+    outflow_edge::Vector{EdgeMetadata}
     active::BitVector
     realized_bmi::Vector{Float64}
     demand::Matrix{Float64}
@@ -632,7 +638,7 @@ struct Parameters{T, C1, C2, V1, V2, V3}
         @NamedTuple{
             node_ids::Dict{Int32, Set{NodeID}},
             edges_source::Dict{Int32, Set{EdgeMetadata}},
-            flow_dict::Dict{Tuple{NodeID, NodeID}, Int},
+            flow_dict::Dict{Tuple{NodeID, NodeID}, Int32},
             flow::T,
             flow_prev::Vector{Float64},
             flow_integrated::Vector{Float64},

core/src/read.jl

@@ -246,8 +246,8 @@ function LinearResistance(db::DB, config::Config, graph::MetaGraph)::LinearResis
 
     return LinearResistance(
         node_id,
-        inflow_id.(Ref(graph), node_id),
-        outflow_id.(Ref(graph), node_id),
+        inflow_edge.(Ref(graph), node_id),
+        outflow_edge.(Ref(graph), node_id),
         BitVector(parsed_parameters.active),
         parsed_parameters.resistance,
         parsed_parameters.max_flow_rate,
@@ -328,16 +328,21 @@ function TabulatedRatingCurve(
 
     return TabulatedRatingCurve(
         node_ids,
-        inflow_id.(Ref(graph), node_ids),
-        [collect(outflow_ids(graph, id)) for id in node_ids],
+        inflow_edge.(Ref(graph), node_ids),
+        outflow_edges.(Ref(graph), node_ids),
         active,
         interpolations,
         time,
         control_mapping,
     )
 end
 
-function ManningResistance(db::DB, config::Config, graph::MetaGraph)::ManningResistance
+function ManningResistance(
+    db::DB,
+    config::Config,
+    graph::MetaGraph,
+    basin::Basin,
+)::ManningResistance
     static = load_structvector(db, config, ManningResistanceStaticV1)
     parsed_parameters, valid =
         parse_static_and_time(db, config, "ManningResistance"; static)
@@ -347,16 +352,20 @@ function ManningResistance(db::DB, config::Config, graph::MetaGraph)::ManningRes
     end
 
     node_id = NodeID.(NodeType.ManningResistance, parsed_parameters.node_id)
+    upstream_bottom = basin_bottom.(Ref(basin), inflow_id.(Ref(graph), node_id))
+    downstream_bottom = basin_bottom.(Ref(basin), outflow_id.(Ref(graph), node_id))
 
     return ManningResistance(
         node_id,
-        inflow_id.(Ref(graph), node_id),
-        outflow_id.(Ref(graph), node_id),
+        inflow_edge.(Ref(graph), node_id),
+        outflow_edge.(Ref(graph), node_id),
         BitVector(parsed_parameters.active),
         parsed_parameters.length,
         parsed_parameters.manning_n,
         parsed_parameters.profile_width,
         parsed_parameters.profile_slope,
+        [bottom[2] for bottom in upstream_bottom],
+        [bottom[2] for bottom in downstream_bottom],
         parsed_parameters.control_mapping,
     )
 end
@@ -373,8 +382,8 @@ function FractionalFlow(db::DB, config::Config, graph::MetaGraph)::FractionalFlo
 
     return FractionalFlow(
         node_id,
-        inflow_id.(Ref(graph), node_id),
-        outflow_id.(Ref(graph), node_id),
+        inflow_edge.(Ref(graph), node_id),
+        outflow_edge.(Ref(graph), node_id),
         parsed_parameters.fraction,
         parsed_parameters.control_mapping,
     )
@@ -464,8 +473,8 @@ function Pump(db::DB, config::Config, graph::MetaGraph, chunk_sizes::Vector{Int}
 
     return Pump(
         node_id,
-        inflow_id.(Ref(graph), node_id),
-        [collect(outflow_ids(graph, id)) for id in node_id],
+        inflow_edge.(Ref(graph), node_id),
+        outflow_edges.(Ref(graph), node_id),
         BitVector(parsed_parameters.active),
         flow_rate,
         parsed_parameters.min_flow_rate,
@@ -497,8 +506,8 @@ function Outlet(db::DB, config::Config, graph::MetaGraph, chunk_sizes::Vector{In
 
     return Outlet(
         node_id,
-        inflow_id.(Ref(graph), node_id),
-        [collect(outflow_ids(graph, id)) for id in node_id],
+        inflow_edge.(Ref(graph), node_id),
+        outflow_edges.(Ref(graph), node_id),
         BitVector(parsed_parameters.active),
         flow_rate,
         parsed_parameters.min_flow_rate,
@@ -889,8 +898,8 @@ function UserDemand(db::DB, config::Config, graph::MetaGraph)::UserDemand
 
     return UserDemand(
         node_ids,
-        inflow_id.(Ref(graph), node_ids),
-        outflow_id.(Ref(graph), node_ids),
+        inflow_edge.(Ref(graph), node_ids),
+        outflow_edge.(Ref(graph), node_ids),
         active,
         realized_bmi,
         demand,
@@ -1071,8 +1080,11 @@ function Parameters(db::DB, config::Config)::Parameters
         error("Invalid number of connections for certain node types.")
     end
 
+    basin = Basin(db, config, graph, chunk_sizes)
+    set_basin_idxs!(graph, basin)
+
     linear_resistance = LinearResistance(db, config, graph)
-    manning_resistance = ManningResistance(db, config, graph)
+    manning_resistance = ManningResistance(db, config, graph, basin)
     tabulated_rating_curve = TabulatedRatingCurve(db, config, graph)
     fractional_flow = FractionalFlow(db, config, graph)
     level_boundary = LevelBoundary(db, config)
@@ -1086,7 +1098,6 @@ function Parameters(db::DB, config::Config)::Parameters
     level_demand = LevelDemand(db, config)
     flow_demand = FlowDemand(db, config)
 
-    basin = Basin(db, config, graph, chunk_sizes)
     subgrid_level = Subgrid(db, config, basin)
 
     p = Parameters(