Precalculate neighboring edges for LinearResistance

core/src/parameter.jl

@@ -235,8 +235,8 @@ 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_edge: incoming flow edge
-outflow_edges: 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
@@ -254,17 +254,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}

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,
@@ -326,15 +326,10 @@ function TabulatedRatingCurve(
         error("Errors occurred when parsing TabulatedRatingCurve data.")
     end
 
-    inflow_edge = [graph[inflow_id(graph, id), id] for id in node_ids]
-    outflow_edges = [
-        [graph[id, outflow_id] for outflow_id in outflow_ids(graph, id)] for id in node_ids
-    ]
-
     return TabulatedRatingCurve(
         node_ids,
-        inflow_edge,
-        outflow_edges,
+        inflow_edge.(Ref(graph), node_ids),
+        outflow_edges.(Ref(graph), node_ids),
         active,
         interpolations,
         time,

core/src/solve.jl

@@ -345,8 +345,11 @@ function formulate_flow!(
     (; graph) = p
     (; node_id, active, resistance, max_flow_rate) = linear_resistance
     for (i, id) in enumerate(node_id)
-        inflow_id = linear_resistance.inflow_id[i]
-        outflow_id = linear_resistance.outflow_id[i]
+        inflow_edge = linear_resistance.inflow_edge[i]
+        outflow_edge = linear_resistance.outflow_edge[i]
+
+        inflow_id = inflow_edge.edge[1]
+        outflow_id = outflow_edge.edge[2]
 
         if active[i]
             _, h_a = get_level(p, inflow_id, t; storage)
@@ -361,8 +364,8 @@ function formulate_flow!(
                 q *= low_storage_factor(storage, p.basin.node_id, outflow_id, 10.0)
             end
 
-            set_flow!(graph, inflow_id, id, q)
-            set_flow!(graph, id, outflow_id, q)
+            set_flow!(graph, inflow_edge, q)
+            set_flow!(graph, outflow_edge, q)
         end
     end
     return nothing

core/src/util.jl

@@ -734,3 +734,8 @@ has_fractional_flow_outneighbors(graph::MetaGraph, node_id::NodeID)::Bool = any(
     outneighbor_id.type == NodeType.FractionalFlow for
     outneighbor_id in outflow_ids(graph, node_id)
 )
+
+inflow_edge(graph, node_id)::EdgeMetadata = graph[inflow_id(graph, node_id), node_id]
+outflow_edge(graph, node_id)::EdgeMetadata = graph[node_id, outflow_id(graph, node_id)]
+outflow_edges(graph, node_id)::Vector{EdgeMetadata} =
+    [graph[node_id, outflow_id] for outflow_id in outflow_ids(graph, node_id)]