Precalculate neighboring edges for Pump

core/src/parameter.jl

@@ -364,8 +364,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
@@ -375,8 +375,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}

core/src/read.jl

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

core/src/solve.jl

@@ -555,10 +555,10 @@ function formulate_flow!(
 )::Nothing
     (; graph, basin) = p
 
-    for (node_id, inflow_id, outflow_ids, active, flow_rate, is_pid_controlled) in zip(
+    for (node_id, inflow_edge, outflow_edges, active, flow_rate, is_pid_controlled) in zip(
         pump.node_id,
-        pump.inflow_id,
-        pump.outflow_ids,
+        pump.inflow_edge,
+        pump.outflow_edges,
         pump.active,
         get_tmp(pump.flow_rate, storage),
         pump.is_pid_controlled,
@@ -567,13 +567,14 @@ function formulate_flow!(
             continue
         end
 
+        inflow_id = inflow_edge.edge[1]
         factor = low_storage_factor(storage, basin.node_id, inflow_id, 10.0)
         q = flow_rate * factor
 
-        set_flow!(graph, inflow_id, node_id, q)
+        set_flow!(graph, inflow_edge, q)
 
-        for outflow_id in outflow_ids
-            set_flow!(graph, node_id, outflow_id, q)
+        for outflow_edge in outflow_edges
+            set_flow!(graph, outflow_edge, q)
         end
     end
     return nothing