Precalculate neighboring edges for TabulatedRatingCurve

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,13 @@ 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),
+        )
         if haskey(graph, id_src, id_dst)
             errors = true
             @error "Duplicate edge" id_src id_dst
@@ -169,8 +175,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
 

core/src/parameter.jl

@@ -116,13 +116,15 @@ 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)
 """
 struct EdgeMetadata
     id::Int32
+    flow_idx::Int32
     type::EdgeType.T
     subnetwork_id_source::Int32
     edge::Tuple{NodeID, NodeID}
@@ -233,17 +235,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
+outflow_edges: outgoing flow edges
 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}
@@ -632,7 +634,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

@@ -326,10 +326,15 @@ 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_id.(Ref(graph), node_ids),
-        [collect(outflow_ids(graph, id)) for id in node_ids],
+        inflow_edge,
+        outflow_edges,
         active,
         interpolations,
         time,

core/src/solve.jl

@@ -378,11 +378,12 @@ function formulate_flow!(
     t::Number,
 )::Nothing
     (; basin, graph) = p
-    (; node_id, active, tables, inflow_id, outflow_ids) = tabulated_rating_curve
+    (; node_id, active, tables, inflow_edge, outflow_edges) = tabulated_rating_curve
 
     for (i, id) in enumerate(node_id)
-        upstream_basin_id = inflow_id[i]
-        downstream_ids = outflow_ids[i]
+        upstream_edge = inflow_edge[i]
+        downstream_edges = outflow_edges[i]
+        upstream_basin_id = upstream_edge.edge[1]
 
         if active[i]
             factor = low_storage_factor(storage, basin.node_id, upstream_basin_id, 10.0)
@@ -391,9 +392,9 @@ function formulate_flow!(
             q = 0.0
         end
 
-        set_flow!(graph, upstream_basin_id, id, q)
-        for downstream_id in downstream_ids
-            set_flow!(graph, id, downstream_id, q)
+        set_flow!(graph, upstream_edge, q)
+        for downstream_edge in downstream_edges
+            set_flow!(graph, downstream_edge, q)
         end
     end
     return nothing

core/test/validation_test.jl

@@ -86,7 +86,7 @@ end
     graph[NodeID(:Pump, 6)] = NodeMetadata(:pump, 9)
 
     function set_edge_metadata!(id_1, id_2, edge_type)
-        graph[id_1, id_2] = EdgeMetadata(0, edge_type, 0, (id_1, id_2))
+        graph[id_1, id_2] = EdgeMetadata(0, 0, edge_type, 0, (id_1, id_2))
         return nothing
     end
 
@@ -163,7 +163,7 @@ end
     graph[NodeID(:Basin, 7)] = NodeMetadata(:basin, 0)
 
     function set_edge_metadata!(id_1, id_2, edge_type)
-        graph[id_1, id_2] = EdgeMetadata(0, edge_type, 0, (id_1, id_2))
+        graph[id_1, id_2] = EdgeMetadata(0, 0, edge_type, 0, (id_1, id_2))
         return nothing
     end