From 55e598cb794a173b1bc946611161c9e086a736ef Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 7 May 2024 19:56:40 +0200
Subject: [PATCH 01/15] Precalculate neighboring edges for TabulatedRatingCurve

---
 core/src/graph.jl            | 26 ++++++++++++++++++++++----
 core/src/parameter.jl        | 12 +++++++-----
 core/src/read.jl             |  9 +++++++--
 core/src/solve.jl            | 13 +++++++------
 core/test/validation_test.jl |  4 ++--
 5 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/core/src/graph.jl b/core/src/graph.jl
index edcd8f3cd..6dc75af10 100644
--- a/core/src/graph.jl
+++ b/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
 
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 5a3672b17..2dac60203 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -116,6 +116,7 @@ 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)
@@ -123,6 +124,7 @@ 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,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_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
@@ -242,8 +244,8 @@ control_mapping: dictionary from (node_id, control_state) to Q(h) and/or active
 """
 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},
diff --git a/core/src/read.jl b/core/src/read.jl
index 523914a41..adafba008 100644
--- a/core/src/read.jl
+++ b/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,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 6d0b2538b..83c6d3144 100644
--- a/core/src/solve.jl
+++ b/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
diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index daf69bfe0..638721261 100644
--- a/core/test/validation_test.jl
+++ b/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
 

From 6425d618bad8c87d94446cd222d79ce44922d82b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 06:39:51 +0200
Subject: [PATCH 02/15] Precalculate neighboring edges for LinearResistance

---
 core/src/parameter.jl | 12 ++++++------
 core/src/read.jl      | 13 ++++---------
 core/src/solve.jl     | 11 +++++++----
 core/src/util.jl      |  5 +++++
 4 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 2dac60203..d015bfaff 100644
--- a/core/src/parameter.jl
+++ b/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,8 +254,8 @@ 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)`
@@ -263,8 +263,8 @@ control_mapping: dictionary from (node_id, control_state) to resistance and/or a
 """
 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}
diff --git a/core/src/read.jl b/core/src/read.jl
index adafba008..ba14405ac 100644
--- a/core/src/read.jl
+++ b/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,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 83c6d3144..736e574aa 100644
--- a/core/src/solve.jl
+++ b/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
diff --git a/core/src/util.jl b/core/src/util.jl
index d03a65b4f..d963f161e 100644
--- a/core/src/util.jl
+++ b/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)]

From a09d4a57f5a78a77060ec6454969d81bc4d809c6 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 06:47:18 +0200
Subject: [PATCH 03/15] Precalculate neighboring edges for ManningResistance

---
 core/src/parameter.jl |  8 ++++----
 core/src/read.jl      |  4 ++--
 core/src/solve.jl     | 11 +++++++----
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index d015bfaff..5483b2a82 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -275,8 +275,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).
 
@@ -309,8 +309,8 @@ 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}
diff --git a/core/src/read.jl b/core/src/read.jl
index ba14405ac..ecf03b2e2 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -350,8 +350,8 @@ function ManningResistance(db::DB, config::Config, graph::MetaGraph)::ManningRes
 
     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,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 736e574aa..63e14c5dd 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -452,8 +452,11 @@ function formulate_flow!(
     (; node_id, active, length, manning_n, profile_width, profile_slope) =
         manning_resistance
     for (i, id) in enumerate(node_id)
-        inflow_id = manning_resistance.inflow_id[i]
-        outflow_id = manning_resistance.outflow_id[i]
+        inflow_edge = manning_resistance.inflow_edge[i]
+        outflow_edge = manning_resistance.outflow_edge[i]
+
+        inflow_id = inflow_edge.edge[1]
+        outflow_id = outflow_edge.edge[2]
 
         if !active[i]
             continue
@@ -492,8 +495,8 @@ function formulate_flow!(
 
         q = q_sign * A / n * R_h^(2 / 3) * sqrt(Δh / L * 2 / π * atan(k * Δh) + eps)
 
-        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
     return nothing
 end

From 9ce0e282696fc5c1de6a0e508e164205a079a9a1 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 07:00:29 +0200
Subject: [PATCH 04/15] Precalculate neighboring edges for FractionalFlow

---
 core/src/graph.jl     | 14 ++++++++++++--
 core/src/parameter.jl |  8 ++++----
 core/src/read.jl      |  4 ++--
 core/src/solve.jl     | 12 ++++++------
 4 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/core/src/graph.jl b/core/src/graph.jl
index 6dc75af10..dfd093452 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -202,9 +202,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
 
 """
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 5483b2a82..2bfcc9994 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -327,15 +327,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
diff --git a/core/src/read.jl b/core/src/read.jl
index ecf03b2e2..a673a8e04 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -373,8 +373,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,
     )
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 63e14c5dd..7047eb2ab 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -509,16 +509,16 @@ function formulate_flow!(
 )::Nothing
     (; graph) = p
 
-    for (node_id, inflow_id, outflow_id, fraction) in zip(
+    for (node_id, inflow_edge, outflow_edge, fraction) in zip(
         fractional_flow.node_id,
-        fractional_flow.inflow_id,
-        fractional_flow.outflow_id,
+        fractional_flow.inflow_edge,
+        fractional_flow.outflow_edge,
         fractional_flow.fraction,
     )
         # overwrite the inflow such that flow is conserved over the FractionalFlow
-        outflow = get_flow(graph, inflow_id, node_id, storage) * fraction
-        set_flow!(graph, inflow_id, node_id, outflow)
-        set_flow!(graph, node_id, outflow_id, outflow)
+        outflow = get_flow(graph, inflow_edge, storage) * fraction
+        set_flow!(graph, inflow_edge, outflow)
+        set_flow!(graph, outflow_edge, outflow)
     end
     return nothing
 end

From 5d0c4dd775177a282af3acec843a01c43d1886cb Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 07:08:21 +0200
Subject: [PATCH 05/15] Precalculate neighboring edges for Pump

---
 core/src/parameter.jl |  8 ++++----
 core/src/read.jl      |  4 ++--
 core/src/solve.jl     | 13 +++++++------
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 2bfcc9994..4ffcdd59e 100644
--- a/core/src/parameter.jl
+++ b/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}
diff --git a/core/src/read.jl b/core/src/read.jl
index a673a8e04..c5f802c6c 100644
--- a/core/src/read.jl
+++ b/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,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 7047eb2ab..b885fe9f2 100644
--- a/core/src/solve.jl
+++ b/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

From d2fef95db4bdd7e1952c05bed868902f9b841001 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 07:19:00 +0200
Subject: [PATCH 06/15] Precalculate neighboring edges for Outlet

---
 core/src/parameter.jl |  8 ++++----
 core/src/read.jl      |  4 ++--
 core/src/solve.jl     | 17 +++++++++--------
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 4ffcdd59e..b7ada4076 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -415,8 +415,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
@@ -426,8 +426,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}
diff --git a/core/src/read.jl b/core/src/read.jl
index c5f802c6c..0a7396e95 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -497,8 +497,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,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index b885fe9f2..6d17057a8 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -590,16 +590,16 @@ function formulate_flow!(
 
     for (
         node_id,
-        inflow_id,
-        outflow_ids,
+        inflow_edge,
+        outflow_edges,
         active,
         flow_rate,
         is_pid_controlled,
         min_crest_level,
     ) in zip(
         outlet.node_id,
-        outlet.inflow_id,
-        outlet.outflow_ids,
+        outlet.inflow_edge,
+        outlet.outflow_edges,
         outlet.active,
         get_tmp(outlet.flow_rate, storage),
         outlet.is_pid_controlled,
@@ -609,12 +609,13 @@ function formulate_flow!(
             continue
         end
 
+        inflow_id = inflow_edge.edge[1]
         q = flow_rate
         q *= low_storage_factor(storage, basin.node_id, inflow_id, 10.0)
 
         # No flow of outlet if source level is lower than target level
         # TODO support multiple outflows to FractionalFlow, or refactor FractionalFlow
-        outflow_id = only(outflow_ids)
+        outflow_id = only(outflow_edges).edge[2]
         _, src_level = get_level(p, inflow_id, t; storage)
         _, dst_level = get_level(p, outflow_id, t; storage)
 
@@ -628,10 +629,10 @@ function formulate_flow!(
             q *= reduction_factor(src_level - min_crest_level, 0.1)
         end
 
-        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

From ad544bb3c6d3f8aa3f0a5f5ec50b663e6928632f Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 07:30:36 +0200
Subject: [PATCH 07/15] Precalculate neighboring edges for UserDemand

---
 core/src/parameter.jl |  8 ++++----
 core/src/read.jl      |  4 ++--
 core/src/solve.jl     | 13 +++++++------
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index b7ada4076..2dba2eff1 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -530,8 +530,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
@@ -547,8 +547,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}
diff --git a/core/src/read.jl b/core/src/read.jl
index 0a7396e95..05b7221b6 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -889,8 +889,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,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 6d17057a8..a71647c54 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -271,8 +271,8 @@ function formulate_flow!(
 
     for (
         node_id,
-        inflow_id,
-        outflow_id,
+        inflow_edge,
+        outflow_edge,
         active,
         demand_itp,
         demand,
@@ -282,8 +282,8 @@ function formulate_flow!(
         demand_from_timeseries,
     ) in zip(
         user_demand.node_id,
-        user_demand.inflow_id,
-        user_demand.outflow_id,
+        user_demand.inflow_edge,
+        user_demand.outflow_edge,
         user_demand.active,
         user_demand.demand_itp,
         # TODO permute these so the nodes are the last dimension, for performance
@@ -315,6 +315,7 @@ function formulate_flow!(
         end
 
         # Smoothly let abstraction go to 0 as the source basin dries out
+        inflow_id = inflow_edge.edge[1]
         factor_basin = low_storage_factor(storage, basin.node_id, inflow_id, 10.0)
         q *= factor_basin
 
@@ -325,10 +326,10 @@ function formulate_flow!(
         factor_level = reduction_factor(Δsource_level, 0.1)
         q *= factor_level
 
-        set_flow!(graph, inflow_id, node_id, q)
+        set_flow!(graph, inflow_edge, q)
 
         # Return flow is immediate
-        set_flow!(graph, node_id, outflow_id, q * return_factor)
+        set_flow!(graph, outflow_edge, q * return_factor)
     end
     return nothing
 end

From 715fffd7947f24ca7d7dd680f6179e583930f759 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 07:54:04 +0200
Subject: [PATCH 08/15] Precompute basin bottoms for Manningresistance

---
 core/src/parameter.jl |  2 ++
 core/src/read.jl      | 15 ++++++++++++---
 core/src/solve.jl     | 16 ++++++++++++----
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 2dba2eff1..33ef2f03a 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -316,6 +316,8 @@ struct ManningResistance <: AbstractParameterNode
     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
 
diff --git a/core/src/read.jl b/core/src/read.jl
index 05b7221b6..5208c16d9 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -337,7 +337,12 @@ function TabulatedRatingCurve(
     )
 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,6 +352,8 @@ 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,
@@ -357,6 +364,8 @@ function ManningResistance(db::DB, config::Config, graph::MetaGraph)::ManningRes
         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
@@ -1071,8 +1080,9 @@ function Parameters(db::DB, config::Config)::Parameters
         error("Invalid number of connections for certain node types.")
     end
 
+    basin = Basin(db, config, graph, chunk_sizes)
     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 +1096,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(
diff --git a/core/src/solve.jl b/core/src/solve.jl
index a71647c54..1f1bae89a 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -450,8 +450,16 @@ function formulate_flow!(
     t::Number,
 )::Nothing
     (; basin, graph) = p
-    (; node_id, active, length, manning_n, profile_width, profile_slope) =
-        manning_resistance
+    (;
+        node_id,
+        active,
+        length,
+        manning_n,
+        profile_width,
+        profile_slope,
+        upstream_bottom,
+        downstream_bottom,
+    ) = manning_resistance
     for (i, id) in enumerate(node_id)
         inflow_edge = manning_resistance.inflow_edge[i]
         outflow_edge = manning_resistance.outflow_edge[i]
@@ -465,8 +473,8 @@ function formulate_flow!(
 
         _, h_a = get_level(p, inflow_id, t; storage)
         _, h_b = get_level(p, outflow_id, t; storage)
-        _, bottom_a = basin_bottom(basin, inflow_id)
-        _, bottom_b = basin_bottom(basin, outflow_id)
+        bottom_a = upstream_bottom[i]
+        bottom_b = downstream_bottom[i]
         slope = profile_slope[i]
         width = profile_width[i]
         n = manning_n[i]

From da045ce7262e787aac31875cc4be5c1c95d2fc36 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 09:03:31 +0200
Subject: [PATCH 09/15] Avoid lookups in formulate_du!

---
 core/src/graph.jl     |  1 +
 core/src/parameter.jl |  2 ++
 core/src/read.jl      |  2 ++
 core/src/solve.jl     | 16 +++++++++++-----
 core/src/util.jl      | 11 +++++++++++
 5 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/core/src/graph.jl b/core/src/graph.jl
index dfd093452..c6f0e6d29 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -71,6 +71,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
             edge_type,
             subnetwork_id,
             (id_src, id_dst),
+            (0, 0),
         )
         if haskey(graph, id_src, id_dst)
             errors = true
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 33ef2f03a..4b03b690d 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -121,6 +121,7 @@ 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
@@ -128,6 +129,7 @@ struct EdgeMetadata
     type::EdgeType.T
     subnetwork_id_source::Int32
     edge::Tuple{NodeID, NodeID}
+    basin_idxs::Tuple{Int32, Int32}
 end
 
 abstract type AbstractParameterNode end
diff --git a/core/src/read.jl b/core/src/read.jl
index 5208c16d9..79904426b 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -1098,6 +1098,8 @@ function Parameters(db::DB, config::Config)::Parameters
 
     subgrid_level = Subgrid(db, config, basin)
 
+    set_basin_idxs!(graph, basin)
+
     p = Parameters(
         config.starttime,
         graph,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 1f1bae89a..8abf0a1ac 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -656,12 +656,18 @@ function formulate_du!(
     # loop over basins
     # subtract all outgoing flows
     # add all ingoing flows
-    for (i, basin_id) in enumerate(basin.node_id)
-        for inflow_id in basin.inflow_ids[i]
-            du[i] += get_flow(graph, inflow_id, basin_id, storage)
+    for edge_metadata in values(graph.edge_data)
+        (; type, edge, basin_idxs) = edge_metadata
+        if type !== EdgeType.flow
+            continue
         end
-        for outflow_id in basin.outflow_ids[i]
-            du[i] -= get_flow(graph, basin_id, outflow_id, storage)
+        q = get_flow(graph, edge_metadata, storage)
+        from_id, to_id = edge
+
+        if from_id.type == NodeType.Basin
+            du[basin_idxs[1]] -= q
+        elseif to_id.type == NodeType.Basin
+            du[basin_idxs[2]] += q
         end
     end
     return nothing
diff --git a/core/src/util.jl b/core/src/util.jl
index d963f161e..0e57d5ed9 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -739,3 +739,14 @@ inflow_edge(graph, node_id)::EdgeMetadata = graph[inflow_id(graph, node_id), nod
 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)]
+
+function set_basin_idxs!(graph::MetaGraph, basin::Basin)::Nothing
+    for edge_metadata in values(graph.edge_data)
+        (; edge) = edge_metadata
+        id_src, id_dst = edge
+        edge_metadata = @set edge_metadata.basin_idxs =
+            (id_index(basin.node_id, id_src)[2], id_index(basin.node_id, id_dst)[2])
+        graph[edge...] = edge_metadata
+    end
+    return nothing
+end

From 30cbbf43b7353895276c69ac95ed0956923aea34 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 09:11:51 +0200
Subject: [PATCH 10/15] Fix tests

---
 core/test/validation_test.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index 638721261..2f181e5fc 100644
--- a/core/test/validation_test.jl
+++ b/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, 0, edge_type, 0, (id_1, id_2))
+        graph[id_1, id_2] = EdgeMetadata(0, 0, edge_type, 0, (id_1, id_2), (0, 0))
         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, 0, edge_type, 0, (id_1, id_2))
+        graph[id_1, id_2] = EdgeMetadata(0, 0, edge_type, 0, (id_1, id_2), (0, 0))
         return nothing
     end
 

From acf434fc7f4c2b1aee48f57af2ac7ac8870c250a Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 10:04:09 +0200
Subject: [PATCH 11/15] Avoid lookups in get_level for UserDemand,
 ManningResistance

---
 core/src/read.jl  |  4 ++--
 core/src/solve.jl | 12 ++++++------
 core/src/util.jl  | 16 ++++++++++++----
 3 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/core/src/read.jl b/core/src/read.jl
index 79904426b..2c7b2df43 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -1081,6 +1081,8 @@ function Parameters(db::DB, config::Config)::Parameters
     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, basin)
     tabulated_rating_curve = TabulatedRatingCurve(db, config, graph)
@@ -1098,8 +1100,6 @@ function Parameters(db::DB, config::Config)::Parameters
 
     subgrid_level = Subgrid(db, config, basin)
 
-    set_basin_idxs!(graph, basin)
-
     p = Parameters(
         config.starttime,
         graph,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 8abf0a1ac..3844ea325 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -25,7 +25,7 @@ function water_balance!(
     formulate_flows!(p, storage, t)
 
     # Now formulate du
-    formulate_du!(du, graph, basin, storage)
+    formulate_du!(du, graph, storage)
 
     # PID control (changes the du of PID controlled basins)
     continuous_control!(u, du, pid_control, p, integral, t)
@@ -321,7 +321,7 @@ function formulate_flow!(
 
         # Smoothly let abstraction go to 0 as the source basin
         # level reaches its minimum level
-        _, source_level = get_level(p, inflow_id, t; storage)
+        source_level = get_level(inflow_edge.basin_idxs[1], basin; storage)
         Δsource_level = source_level - min_level
         factor_level = reduction_factor(Δsource_level, 0.1)
         q *= factor_level
@@ -471,8 +471,8 @@ function formulate_flow!(
             continue
         end
 
-        _, h_a = get_level(p, inflow_id, t; storage)
-        _, h_b = get_level(p, outflow_id, t; storage)
+        h_a = get_level(inflow_edge.basin_idxs[1], basin; storage)
+        h_b = get_level(outflow_edge.basin_idxs[2], basin; storage)
         bottom_a = upstream_bottom[i]
         bottom_b = downstream_bottom[i]
         slope = profile_slope[i]
@@ -650,7 +650,6 @@ end
 function formulate_du!(
     du::ComponentVector,
     graph::MetaGraph,
-    basin::Basin,
     storage::AbstractVector,
 )::Nothing
     # loop over basins
@@ -661,12 +660,13 @@ function formulate_du!(
         if type !== EdgeType.flow
             continue
         end
-        q = get_flow(graph, edge_metadata, storage)
         from_id, to_id = edge
 
         if from_id.type == NodeType.Basin
+            q = get_flow(graph, edge_metadata, storage)
             du[basin_idxs[1]] -= q
         elseif to_id.type == NodeType.Basin
+            q = get_flow(graph, edge_metadata, storage)
             du[basin_idxs[2]] += q
         end
     end
diff --git a/core/src/util.jl b/core/src/util.jl
index 0e57d5ed9..0b918e0dd 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -380,8 +380,7 @@ function get_level(
     (; basin, level_boundary) = p
     if node_id.type == NodeType.Basin
         _, i = id_index(basin.node_id, node_id)
-        current_level = get_tmp(basin.current_level, storage)
-        return true, current_level[i]
+        return true, get_level(i, basin; storage)
     elseif node_id.type == NodeType.LevelBoundary
         i = findsorted(level_boundary.node_id, node_id)
         return true, level_boundary.level[i](t)
@@ -390,6 +389,14 @@ function get_level(
     end
 end
 
+function get_level(
+    i::Integer,
+    basin::Basin;
+    storage::Union{AbstractArray, Number} = 0,
+)::Number
+    return get_tmp(basin.current_level, storage)[i]
+end
+
 "Get the index of an ID in a set of indices."
 function id_index(ids::Indices{NodeID}, id::NodeID)::Tuple{Bool, Int}
     # We avoid creating Dictionary here since it converts the values to a Vector,
@@ -741,12 +748,13 @@ outflow_edges(graph, node_id)::Vector{EdgeMetadata} =
     [graph[node_id, outflow_id] for outflow_id in outflow_ids(graph, node_id)]
 
 function set_basin_idxs!(graph::MetaGraph, basin::Basin)::Nothing
-    for edge_metadata in values(graph.edge_data)
-        (; edge) = edge_metadata
+    for (edge, edge_metadata) in graph.edge_data
         id_src, id_dst = edge
         edge_metadata = @set edge_metadata.basin_idxs =
             (id_index(basin.node_id, id_src)[2], id_index(basin.node_id, id_dst)[2])
         graph[edge...] = edge_metadata
+        if edge[2] == NodeID(:ManningResistance, 8)
+        end
     end
     return nothing
 end

From 57b701595b53ff43a215283e3d9ff44fbdfe50e3 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 11:44:52 +0200
Subject: [PATCH 12/15] Add basin index test

---
 core/test/validation_test.jl | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index 2f181e5fc..4cbf242da 100644
--- a/core/test/validation_test.jl
+++ b/core/test/validation_test.jl
@@ -452,3 +452,22 @@ end
         dt,
     )
 end
+
+@testitem "basin indices" begin
+    using Ribasim: NodeType
+
+    toml_path = normpath(@__DIR__, "../../generated_testmodels/basic/ribasim.toml")
+    @test ispath(toml_path)
+
+    model = Ribasim.Model(toml_path)
+    (; graph, basin) = model.integrator.p
+    for edge_metadata in values(graph.edge_data)
+        (; edge, basin_idxs) = edge_metadata
+        id_src, id_dst = edge
+        if id_src.type == NodeType.Basin
+            @test id_src == basin.node_id.values[basin_idxs[1]]
+        elseif id_dst.type == NodeType.Basin
+            @test id_dst == basin.node_id.values[basin_idxs[2]]
+        end
+    end
+end

From e88f0e0e62c9317a0a9e47308f105bab1fe8acd7 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 8 May 2024 15:26:22 +0200
Subject: [PATCH 13/15] Don't lookup basin index for bottoms for UserDemand,
 Linearresistance, TRC, Manning, Outlet

---
 core/src/solve.jl | 34 ++++++++++++++++++++++------------
 core/src/util.jl  | 14 ++++----------
 2 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/core/src/solve.jl b/core/src/solve.jl
index 3844ea325..856dfa969 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -148,8 +148,8 @@ function continuous_control!(
             src_id = inflow_id(graph, controlled_node_id)
             dst_id = outflow_id(graph, controlled_node_id)
 
-            has_src_level, src_level = get_level(p, src_id, t; storage)
-            has_dst_level, dst_level = get_level(p, dst_id, t; storage)
+            has_src_level, src_level = get_level(p, src_id, 0, t; storage)
+            has_dst_level, dst_level = get_level(p, dst_id, 0, t; storage)
 
             factor_outlet = 1.0
 
@@ -321,7 +321,7 @@ function formulate_flow!(
 
         # Smoothly let abstraction go to 0 as the source basin
         # level reaches its minimum level
-        source_level = get_level(inflow_edge.basin_idxs[1], basin; storage)
+        _, source_level = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
         Δsource_level = source_level - min_level
         factor_level = reduction_factor(Δsource_level, 0.1)
         q *= factor_level
@@ -353,8 +353,8 @@ function formulate_flow!(
         outflow_id = outflow_edge.edge[2]
 
         if active[i]
-            _, h_a = get_level(p, inflow_id, t; storage)
-            _, h_b = get_level(p, outflow_id, t; storage)
+            _, h_a = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
+            _, h_b = get_level(p, outflow_id, outflow_edge.basin_idxs[2], t; storage)
             q_unlimited = (h_a - h_b) / resistance[i]
             q = clamp(q_unlimited, -max_flow_rate[i], max_flow_rate[i])
 
@@ -391,7 +391,16 @@ function formulate_flow!(
 
         if active[i]
             factor = low_storage_factor(storage, basin.node_id, upstream_basin_id, 10.0)
-            q = factor * tables[i](get_level(p, upstream_basin_id, t; storage)[2])
+            q =
+                factor * tables[i](
+                    get_level(
+                        p,
+                        upstream_basin_id,
+                        upstream_edge.basin_idxs[1],
+                        t;
+                        storage,
+                    )[2],
+                )
         else
             q = 0.0
         end
@@ -449,7 +458,7 @@ function formulate_flow!(
     storage::AbstractVector,
     t::Number,
 )::Nothing
-    (; basin, graph) = p
+    (; graph) = p
     (;
         node_id,
         active,
@@ -471,8 +480,8 @@ function formulate_flow!(
             continue
         end
 
-        h_a = get_level(inflow_edge.basin_idxs[1], basin; storage)
-        h_b = get_level(outflow_edge.basin_idxs[2], basin; storage)
+        _, h_a = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
+        _, h_b = get_level(p, outflow_id, outflow_edge.basin_idxs[2], t; storage)
         bottom_a = upstream_bottom[i]
         bottom_b = downstream_bottom[i]
         slope = profile_slope[i]
@@ -624,9 +633,10 @@ function formulate_flow!(
 
         # No flow of outlet if source level is lower than target level
         # TODO support multiple outflows to FractionalFlow, or refactor FractionalFlow
-        outflow_id = only(outflow_edges).edge[2]
-        _, src_level = get_level(p, inflow_id, t; storage)
-        _, dst_level = get_level(p, outflow_id, t; storage)
+        outflow_edge = only(outflow_edges)
+        outflow_id = outflow_edge.edge[2]
+        _, src_level = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
+        _, dst_level = get_level(p, outflow_id, outflow_edge.basin_idxs[2], t; storage)
 
         if src_level !== nothing && dst_level !== nothing
             Δlevel = src_level - dst_level
diff --git a/core/src/util.jl b/core/src/util.jl
index 0b918e0dd..642fcf9f7 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -374,13 +374,15 @@ storage: tells ForwardDiff whether this call is for differentiation or not
 function get_level(
     p::Parameters,
     node_id::NodeID,
+    i::Integer,
     t::Number;
     storage::Union{AbstractArray, Number} = 0,
 )::Tuple{Bool, Number}
     (; basin, level_boundary) = p
     if node_id.type == NodeType.Basin
-        _, i = id_index(basin.node_id, node_id)
-        return true, get_level(i, basin; storage)
+        i = iszero(i) ? id_index(basin.node_id, node_id)[2] : i
+        current_level = get_tmp(basin.current_level, storage)
+        return true, current_level[i]
     elseif node_id.type == NodeType.LevelBoundary
         i = findsorted(level_boundary.node_id, node_id)
         return true, level_boundary.level[i](t)
@@ -389,14 +391,6 @@ function get_level(
     end
 end
 
-function get_level(
-    i::Integer,
-    basin::Basin;
-    storage::Union{AbstractArray, Number} = 0,
-)::Number
-    return get_tmp(basin.current_level, storage)[i]
-end
-
 "Get the index of an ID in a set of indices."
 function id_index(ids::Indices{NodeID}, id::NodeID)::Tuple{Bool, Int}
     # We avoid creating Dictionary here since it converts the values to a Vector,

From f18cd51ce847ad3bfd1f1ca52f51b998fb97f3b3 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 13 May 2024 12:57:54 +0200
Subject: [PATCH 14/15] Comments adressed + avoid lookups for
 low_storage_factor and other small performance improvements

---
 core/src/callback.jl         |  4 +-
 core/src/graph.jl            | 43 ++++++++++-------
 core/src/parameter.jl        | 42 +++++++++-------
 core/src/solve.jl            | 72 ++++++++++++++--------------
 core/src/util.jl             | 93 +++++++++++++++++++-----------------
 core/test/utils_test.jl      | 54 +++++----------------
 core/test/validation_test.jl | 10 ++--
 7 files changed, 154 insertions(+), 164 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 74ae2b345..908e21b03 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -125,9 +125,7 @@ function integrate_flows!(u, t, integrator)::Nothing
         else
             # Horizontal flows
             value[] +=
-                0.5 *
-                (get_flow(graph, edge..., 0) + get_flow(graph, edge..., 0; prev = true)) *
-                dt
+                0.5 * (get_flow(graph, edge..., 0) + get_flow_prev(graph, edge..., 0)) * dt
         end
     end
 
diff --git a/core/src/graph.jl b/core/src/graph.jl
index c6f0e6d29..8351dd92d 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -18,7 +18,8 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
     node_ids = Dict{Int32, Set{NodeID}}()
     # Source edges per subnetwork
     edges_source = Dict{Int32, Set{EdgeMetadata}}()
-    # The number of flow edges
+    # The flow counter gives a unique consecutive id to the
+    # flow edges to index the flow vectors
     flow_counter = 0
     # Dictionary from flow edge to index in flow vector
     flow_dict = Dict{Tuple{NodeID, NodeID}, Int32}()
@@ -71,7 +72,8 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
             edge_type,
             subnetwork_id,
             (id_src, id_dst),
-            (0, 0),
+            -1,
+            -1,
         )
         if haskey(graph, id_src, id_dst)
             errors = true
@@ -196,26 +198,35 @@ end
 """
 Get the flow over the given edge (val is needed for get_tmp from ForwardDiff.jl).
 """
-function get_flow(
-    graph::MetaGraph,
-    id_src::NodeID,
-    id_dst::NodeID,
-    val;
-    prev::Bool = false,
-)::Number
+function get_flow(graph::MetaGraph, id_src::NodeID, id_dst::NodeID, val)::Number
+    (; flow_dict) = graph[]
+    flow_idx = flow_dict[id_src, id_dst]
+    return get_flow(graph, flow_idx, val)
+end
+
+function get_flow(graph, edge_metadata::EdgeMetadata, val)::Number
+    return get_flow(graph, edge_metadata.flow_idx, val)
+end
+
+function get_flow(graph::MetaGraph, flow_idx::Integer, val)
+    return get_tmp(graph[].flow, val)[flow_idx]
+end
+
+function get_flow_prev(graph, id_src::NodeID, id_dst::NodeID, val)::Number
+    # Note: Can be removed after https://github.com/Deltares/Ribasim/pull/1444
     (; flow_dict) = graph[]
     flow_idx = flow_dict[id_src, id_dst]
-    return get_flow(graph, flow_idx, val; prev)
+    return get_flow(graph, flow_idx, val)
 end
 
-function get_flow(graph, edge_metadata::EdgeMetadata, val; prev::Bool = false)::Number
-    return get_flow(graph, edge_metadata.flow_idx, val; prev)
+function get_flow_prev(graph, edge_metadata::EdgeMetadata, val)::Number
+    # Note: Can be removed after https://github.com/Deltares/Ribasim/pull/1444
+    return get_flow_prev(graph, edge_metadata.flow_idx, val)
 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_idx]
+function get_flow_prev(graph::MetaGraph, flow_idx::Integer, val)
+    # Note: Can be removed after https://github.com/Deltares/Ribasim/pull/1444
+    return get_tmp(graph[].flow_prev, val)[flow_idx]
 end
 
 """
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 4b03b690d..ccb017ca2 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -121,7 +121,8 @@ 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)
+basin_idx_src: Basin index of source node (0 if not a basin)
+basin_idx_dst: Basin index of destination node (0 if not a basin)
 """
 struct EdgeMetadata
     id::Int32
@@ -129,7 +130,8 @@ struct EdgeMetadata
     type::EdgeType.T
     subnetwork_id_source::Int32
     edge::Tuple{NodeID, NodeID}
-    basin_idxs::Tuple{Int32, Int32}
+    basin_idx_src::Int32
+    basin_idx_dst::Int32
 end
 
 abstract type AbstractParameterNode end
@@ -238,7 +240,9 @@ 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 metadata
+    The ID of the destination node is always the ID of the TabulatedRatingCurve node
 outflow_edges: outgoing flow edges metadata
+    The ID of the source node is always the ID of the TabulatedRatingCurve node
 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
@@ -257,7 +261,9 @@ end
 """
 node_id: node ID of the LinearResistance node
 inflow_edge: incoming flow edge metadata
+    The ID of the destination node is always the ID of the LinearResistance node
 outflow_edge: outgoing flow edge metadata
+    The ID of the source node is always the ID of the LinearResistance node
 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)`
@@ -278,7 +284,9 @@ This is a simple Manning-Gauckler reach connection.
 
 node_id: node ID of the ManningResistance node
 inflow_edge: incoming flow edge metadata
+    The ID of the destination node is always the ID of the ManningResistance node
 outflow_edge: outgoing flow edge metadata
+    The ID of the source node is always the ID of the ManningResistance node
 length: reach length
 manning_n: roughness; Manning's n in (SI units).
 
@@ -324,15 +332,11 @@ struct ManningResistance <: AbstractParameterNode
 end
 
 """
-Requirements:
-
-* from: must be (TabulatedRatingCurve,) node
-* to: must be (Basin,) node
-* fraction must be positive.
-
-node_id: node ID of the TabulatedRatingCurve node
+node_id: node ID of the FractionalFlow node
 inflow_edge: incoming flow edge metadata
+    The ID of the destination node is always the ID of the FractionalFlow node
 outflow_edge: outgoing flow edge metadata
+    The ID of the source node is always the ID of the FractionalFlow node
 fraction: The fraction in [0,1] of flow the node lets through
 control_mapping: dictionary from (node_id, control_state) to fraction
 """
@@ -369,7 +373,9 @@ end
 """
 node_id: node ID of the Pump node
 inflow_edge: incoming flow edge metadata
+    The ID of the destination node is always the ID of the Pump node
 outflow_edges: outgoing flow edges metadata
+    The ID of the source node is always the ID of the Pump node
 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
@@ -390,8 +396,8 @@ struct Pump{T} <: AbstractParameterNode
 
     function Pump(
         node_id,
-        inflow_id,
-        outflow_ids,
+        inflow_edge,
+        outflow_edges,
         active,
         flow_rate::T,
         min_flow_rate,
@@ -402,8 +408,8 @@ struct Pump{T} <: AbstractParameterNode
         if valid_flow_rates(node_id, get_tmp(flow_rate, 0), control_mapping)
             return new{T}(
                 node_id,
-                inflow_id,
-                outflow_ids,
+                inflow_edge,
+                outflow_edges,
                 active,
                 flow_rate,
                 min_flow_rate,
@@ -419,8 +425,10 @@ end
 
 """
 node_id: node ID of the Outlet node
-inflow_edge: incoming flow edge metadata
-outflow_edges: outgoing flow edges metadata
+inflow_edge: incoming flow edge metadata.
+    The ID of the destination node is always the ID of the Outlet node
+outflow_edges: outgoing flow edges metadata.
+    The ID of the source node is always the ID of the Outlet node
 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
@@ -534,8 +542,10 @@ end
 
 """
 node_id: node ID of the UserDemand node
-inflow_edge: incoming flow edge metadata
+inflow_edge: incoming flow edge
+    The ID of the destination node is always the ID of the UserDemand node
 outflow_edge: outgoing flow edge metadata
+    The ID of the source node is always the ID of the UserDemand node
 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
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 856dfa969..a0334bee0 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -40,7 +40,10 @@ function set_current_basin_properties!(basin::Basin, storage::AbstractVector)::N
 
     for i in eachindex(storage)
         s = storage[i]
-        current_area[i], current_level[i] = get_area_and_level(basin, i, s)
+        area, level = get_area_and_level(basin, i, s)
+
+        current_area[i] = area
+        current_level[i] = level
     end
 end
 
@@ -148,8 +151,11 @@ function continuous_control!(
             src_id = inflow_id(graph, controlled_node_id)
             dst_id = outflow_id(graph, controlled_node_id)
 
-            has_src_level, src_level = get_level(p, src_id, 0, t; storage)
-            has_dst_level, dst_level = get_level(p, dst_id, 0, t; storage)
+            inflow_edge = graph[src_id, controlled_node_id]
+            outflow_edge = graph[controlled_node_id, dst_id]
+
+            has_src_level, src_level = get_level(p, inflow_edge, src_id, t; storage)
+            has_dst_level, dst_level = get_level(p, outflow_edge, dst_id, t; storage)
 
             factor_outlet = 1.0
 
@@ -173,7 +179,8 @@ function continuous_control!(
         end
 
         id_inflow = inflow_id(graph, controlled_node_id)
-        factor_basin = low_storage_factor(storage, basin.node_id, id_inflow, 10.0)
+        inflow_edge = graph[id_inflow, controlled_node_id]
+        factor_basin = low_storage_factor(storage, inflow_edge, id_inflow, 10.0)
 
         factor = factor_basin * factor_outlet
         flow_rate = 0.0
@@ -267,7 +274,7 @@ function formulate_flow!(
     storage::AbstractVector,
     t::Number,
 )::Nothing
-    (; graph, basin, allocation) = p
+    (; graph, allocation) = p
 
     for (
         node_id,
@@ -316,12 +323,12 @@ function formulate_flow!(
 
         # Smoothly let abstraction go to 0 as the source basin dries out
         inflow_id = inflow_edge.edge[1]
-        factor_basin = low_storage_factor(storage, basin.node_id, inflow_id, 10.0)
+        factor_basin = low_storage_factor(storage, inflow_edge, inflow_id, 10.0)
         q *= factor_basin
 
         # Smoothly let abstraction go to 0 as the source basin
         # level reaches its minimum level
-        _, source_level = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
+        _, source_level = get_level(p, inflow_edge, inflow_id, t; storage)
         Δsource_level = source_level - min_level
         factor_level = reduction_factor(Δsource_level, 0.1)
         q *= factor_level
@@ -353,16 +360,16 @@ function formulate_flow!(
         outflow_id = outflow_edge.edge[2]
 
         if active[i]
-            _, h_a = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
-            _, h_b = get_level(p, outflow_id, outflow_edge.basin_idxs[2], t; storage)
+            _, h_a = get_level(p, inflow_edge, inflow_id, t; storage)
+            _, h_b = get_level(p, outflow_edge, outflow_id, t; storage)
             q_unlimited = (h_a - h_b) / resistance[i]
             q = clamp(q_unlimited, -max_flow_rate[i], max_flow_rate[i])
 
             # add reduction_factor on highest level
             if q > 0
-                q *= low_storage_factor(storage, p.basin.node_id, inflow_id, 10.0)
+                q *= low_storage_factor(storage, inflow_edge, inflow_id, 10.0)
             else
-                q *= low_storage_factor(storage, p.basin.node_id, outflow_id, 10.0)
+                q *= low_storage_factor(storage, outflow_edge, outflow_id, 10.0)
             end
 
             set_flow!(graph, inflow_edge, q)
@@ -381,7 +388,7 @@ function formulate_flow!(
     storage::AbstractVector,
     t::Number,
 )::Nothing
-    (; basin, graph) = p
+    (; graph) = p
     (; node_id, active, tables, inflow_edge, outflow_edges) = tabulated_rating_curve
 
     for (i, id) in enumerate(node_id)
@@ -390,17 +397,10 @@ function formulate_flow!(
         upstream_basin_id = upstream_edge.edge[1]
 
         if active[i]
-            factor = low_storage_factor(storage, basin.node_id, upstream_basin_id, 10.0)
+            factor = low_storage_factor(storage, upstream_edge, upstream_basin_id, 10.0)
             q =
-                factor * tables[i](
-                    get_level(
-                        p,
-                        upstream_basin_id,
-                        upstream_edge.basin_idxs[1],
-                        t;
-                        storage,
-                    )[2],
-                )
+                factor *
+                tables[i](get_level(p, upstream_edge, upstream_basin_id, t; storage)[2])
         else
             q = 0.0
         end
@@ -455,9 +455,9 @@ dry.
 function formulate_flow!(
     manning_resistance::ManningResistance,
     p::Parameters,
-    storage::AbstractVector,
+    storage::AbstractVector{T},
     t::Number,
-)::Nothing
+)::Nothing where {T}
     (; graph) = p
     (;
         node_id,
@@ -480,8 +480,8 @@ function formulate_flow!(
             continue
         end
 
-        _, h_a = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
-        _, h_b = get_level(p, outflow_id, outflow_edge.basin_idxs[2], t; storage)
+        _, h_a = get_level(p, inflow_edge, inflow_id, t; storage)
+        _, h_b = get_level(p, outflow_edge, outflow_id, t; storage)
         bottom_a = upstream_bottom[i]
         bottom_b = downstream_bottom[i]
         slope = profile_slope[i]
@@ -506,12 +506,12 @@ function formulate_flow!(
         P_b = width + 2.0 * d_b * slope_unit_length
         R_h_a = A_a / P_a
         R_h_b = A_b / P_b
-        R_h = 0.5 * (R_h_a + R_h_b)
+        R_h::T = 0.5 * (R_h_a + R_h_b)
         k = 1000.0
         # This epsilon makes sure the AD derivative at Δh = 0 does not give NaN
         eps = 1e-200
 
-        q = q_sign * A / n * R_h^(2 / 3) * sqrt(Δh / L * 2 / π * atan(k * Δh) + eps)
+        q = q_sign * A / n * ∛(R_h^2) * sqrt(Δh / L * 2 / π * atan(k * Δh) + eps)
 
         set_flow!(graph, inflow_edge, q)
         set_flow!(graph, outflow_edge, q)
@@ -586,7 +586,7 @@ function formulate_flow!(
         end
 
         inflow_id = inflow_edge.edge[1]
-        factor = low_storage_factor(storage, basin.node_id, inflow_id, 10.0)
+        factor = low_storage_factor(storage, inflow_edge, inflow_id, 10.0)
         q = flow_rate * factor
 
         set_flow!(graph, inflow_edge, q)
@@ -604,7 +604,7 @@ function formulate_flow!(
     storage::AbstractVector,
     t::Number,
 )::Nothing
-    (; graph, basin) = p
+    (; graph) = p
 
     for (
         node_id,
@@ -629,14 +629,14 @@ function formulate_flow!(
 
         inflow_id = inflow_edge.edge[1]
         q = flow_rate
-        q *= low_storage_factor(storage, basin.node_id, inflow_id, 10.0)
+        q *= low_storage_factor(storage, inflow_edge, inflow_id, 10.0)
 
         # No flow of outlet if source level is lower than target level
         # TODO support multiple outflows to FractionalFlow, or refactor FractionalFlow
         outflow_edge = only(outflow_edges)
         outflow_id = outflow_edge.edge[2]
-        _, src_level = get_level(p, inflow_id, inflow_edge.basin_idxs[1], t; storage)
-        _, dst_level = get_level(p, outflow_id, outflow_edge.basin_idxs[2], t; storage)
+        _, src_level = get_level(p, inflow_edge, inflow_id, t; storage)
+        _, dst_level = get_level(p, outflow_edge, outflow_id, t; storage)
 
         if src_level !== nothing && dst_level !== nothing
             Δlevel = src_level - dst_level
@@ -666,7 +666,7 @@ function formulate_du!(
     # subtract all outgoing flows
     # add all ingoing flows
     for edge_metadata in values(graph.edge_data)
-        (; type, edge, basin_idxs) = edge_metadata
+        (; type, edge, basin_idx_src, basin_idx_dst) = edge_metadata
         if type !== EdgeType.flow
             continue
         end
@@ -674,10 +674,10 @@ function formulate_du!(
 
         if from_id.type == NodeType.Basin
             q = get_flow(graph, edge_metadata, storage)
-            du[basin_idxs[1]] -= q
+            du[basin_idx_src] -= q
         elseif to_id.type == NodeType.Basin
             q = get_flow(graph, edge_metadata, storage)
-            du[basin_idxs[2]] += q
+            du[basin_idx_dst] += q
         end
     end
     return nothing
diff --git a/core/src/util.jl b/core/src/util.jl
index 642fcf9f7..4cc6d0f84 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -89,11 +89,7 @@ end
 Compute the area and level of a basin given its storage.
 Also returns darea/dlevel as it is needed for the Jacobian.
 """
-function get_area_and_level(
-    basin::Basin,
-    state_idx::Int,
-    storage::Number,
-)::Tuple{Real, Real}
+function get_area_and_level(basin::Basin, state_idx::Int, storage::T)::Tuple{T, T} where {T}
     storage_discrete = basin.storage[state_idx]
     area_discrete = basin.area[state_idx]
     level_discrete = basin.level[state_idx]
@@ -102,11 +98,16 @@ function get_area_and_level(
 end
 
 function get_area_and_level(
-    storage_discrete::Vector,
-    area_discrete::Vector,
-    level_discrete::Vector,
-    storage::Number,
-)::Tuple{Number, Number}
+    storage_discrete::AbstractVector,
+    area_discrete::AbstractVector,
+    level_discrete::AbstractVector,
+    storage::T,
+)::Tuple{T, T} where {T}
+
+    # Set type of area and level to prevent runtime dispatch
+    area::T = zero(T)
+    level::T = zero(T)
+
     # storage_idx: smallest index such that storage_discrete[storage_idx] >= storage
     storage_idx = searchsortedfirst(storage_discrete, storage)
 
@@ -115,13 +116,6 @@ function get_area_and_level(
         level = level_discrete[1]
         area = area_discrete[1]
 
-        level_lower = level
-        level_higher = level_discrete[2]
-        area_lower = area
-        area_higher = area_discrete[2]
-
-        darea = (area_higher - area_lower) / (level_higher - level_lower)
-
     elseif storage_idx == length(storage_discrete) + 1
         # With a storage above the profile, use a linear extrapolation of area(level)
         # based on the last 2 values.
@@ -132,20 +126,18 @@ function get_area_and_level(
         storage_lower = storage_discrete[end - 1]
         storage_higher = storage_discrete[end]
 
-        area_diff = area_higher - area_lower
-        level_diff = level_higher - level_lower
+        Δarea = area_higher - area_lower
+        Δlevel = level_higher - level_lower
+        Δstorage = storage_higher - storage_lower
 
-        if area_diff ≈ 0
+        if Δarea ≈ 0.0
             # Constant area means linear interpolation of level
-            darea = 0.0
             area = area_lower
-            level =
-                level_higher +
-                level_diff * (storage - storage_higher) / (storage_higher - storage_lower)
+            level = level_higher + Δlevel * (storage - storage_higher) / Δstorage
         else
-            darea = area_diff / level_diff
+            darea = Δarea / Δlevel
             area = sqrt(area_higher^2 + 2 * (storage - storage_higher) * darea)
-            level = level_lower + level_diff * (area - area_lower) / area_diff
+            level = level_lower + Δlevel * (area - area_lower) / Δarea
         end
 
     else
@@ -156,21 +148,19 @@ function get_area_and_level(
         storage_lower = storage_discrete[storage_idx - 1]
         storage_higher = storage_discrete[storage_idx]
 
-        area_diff = area_higher - area_lower
-        level_diff = level_higher - level_lower
+        Δarea = area_higher - area_lower
+        Δlevel = level_higher - level_lower
+        Δstorage = storage_higher - storage_lower
 
-        if area_diff ≈ 0
+        if Δarea ≈ 0.0
             # Constant area means linear interpolation of level
-            darea = 0.0
             area = area_lower
-            level =
-                level_lower +
-                level_diff * (storage - storage_lower) / (storage_higher - storage_lower)
+            level = level_lower + Δlevel * (storage - storage_lower) / Δstorage
 
         else
-            darea = area_diff / level_diff
+            darea = Δarea / Δlevel
             area = sqrt(area_lower^2 + 2 * (storage - storage_lower) * darea)
-            level = level_lower + level_diff * (area - area_lower) / area_diff
+            level = level_lower + Δlevel * (area - area_lower) / Δarea
         end
     end
 
@@ -373,14 +363,16 @@ storage: tells ForwardDiff whether this call is for differentiation or not
 """
 function get_level(
     p::Parameters,
+    edge_metadata::EdgeMetadata,
     node_id::NodeID,
-    i::Integer,
     t::Number;
     storage::Union{AbstractArray, Number} = 0,
 )::Tuple{Bool, Number}
     (; basin, level_boundary) = p
     if node_id.type == NodeType.Basin
-        i = iszero(i) ? id_index(basin.node_id, node_id)[2] : i
+        # The edge metadata is only used to obtain the Basin index
+        # in case node_id is for a Basin
+        i = get_basin_idx(edge_metadata, node_id)
         current_level = get_tmp(basin.current_level, storage)
         return true, current_level[i]
     elseif node_id.type == NodeType.LevelBoundary
@@ -550,13 +542,13 @@ end
 "If id is a Basin with storage below the threshold, return a reduction factor != 1"
 function low_storage_factor(
     storage::AbstractVector{T},
-    basin_ids::Indices{NodeID},
+    edge_metadata::EdgeMetadata,
     id::NodeID,
     threshold::Real,
 )::T where {T <: Real}
-    hasindex, basin_idx = id_index(basin_ids, id)
-    return if hasindex
-        reduction_factor(storage[basin_idx], threshold)
+    if id.type == NodeType.Basin
+        i = get_basin_idx(edge_metadata, id)
+        reduction_factor(storage[i], threshold)
     else
         one(T)
     end
@@ -744,11 +736,22 @@ outflow_edges(graph, node_id)::Vector{EdgeMetadata} =
 function set_basin_idxs!(graph::MetaGraph, basin::Basin)::Nothing
     for (edge, edge_metadata) in graph.edge_data
         id_src, id_dst = edge
-        edge_metadata = @set edge_metadata.basin_idxs =
-            (id_index(basin.node_id, id_src)[2], id_index(basin.node_id, id_dst)[2])
+        edge_metadata =
+            @set edge_metadata.basin_idx_src = id_index(basin.node_id, id_src)[2]
+        edge_metadata =
+            @set edge_metadata.basin_idx_dst = id_index(basin.node_id, id_dst)[2]
         graph[edge...] = edge_metadata
-        if edge[2] == NodeID(:ManningResistance, 8)
-        end
     end
     return nothing
 end
+
+function get_basin_idx(edge_metadata::EdgeMetadata, id::NodeID)::Int32
+    (; edge) = edge_metadata
+    return if edge[1] == id
+        edge_metadata.basin_idx_src
+    elseif edge[2] == id
+        edge_metadata.basin_idx_dst
+    else
+        0
+    end
+end
diff --git a/core/test/utils_test.jl b/core/test/utils_test.jl
index ff2520bca..50bfe9076 100644
--- a/core/test/utils_test.jl
+++ b/core/test/utils_test.jl
@@ -240,49 +240,17 @@ end
 end
 
 @testitem "low_storage_factor" begin
-    using Ribasim: NodeID, low_storage_factor, Indices
-    @test low_storage_factor(
-        [-2.0],
-        Indices(NodeID.(:Basin, [5])),
-        NodeID(:Basin, 5),
-        2.0,
-    ) === 0.0
-    @test low_storage_factor(
-        [0.0f0],
-        Indices(NodeID.(:Basin, [5])),
-        NodeID(:Basin, 5),
-        2.0,
-    ) === 0.0f0
-    @test low_storage_factor(
-        [0.0],
-        Indices(NodeID.(:Basin, [5])),
-        NodeID(:Basin, 5),
-        2.0,
-    ) === 0.0
-    @test low_storage_factor(
-        [1.0f0],
-        Indices(NodeID.(:Basin, [5])),
-        NodeID(:Basin, 5),
-        2.0,
-    ) === 0.5f0
-    @test low_storage_factor(
-        [1.0],
-        Indices(NodeID.(:Basin, [5])),
-        NodeID(:Basin, 5),
-        2.0,
-    ) === 0.5
-    @test low_storage_factor(
-        [3.0f0],
-        Indices(NodeID.(:Basin, [5])),
-        NodeID(:Basin, 5),
-        2.0,
-    ) === 1.0f0
-    @test low_storage_factor(
-        [3.0],
-        Indices(NodeID.(:Basin, [5])),
-        NodeID(:Basin, 5),
-        2.0,
-    ) === 1.0
+    using Ribasim: NodeID, low_storage_factor, EdgeMetadata, EdgeType
+
+    node_id = NodeID(:Basin, 5)
+    edge_metadata = EdgeMetadata(0, 0, EdgeType.flow, 0, (node_id, node_id), 1, 1)
+    @test low_storage_factor([-2.0], edge_metadata, node_id, 2.0) === 0.0
+    @test low_storage_factor([0.0f0], edge_metadata, node_id, 2.0) === 0.0f0
+    @test low_storage_factor([0.0], edge_metadata, node_id, 2.0) === 0.0
+    @test low_storage_factor([1.0f0], edge_metadata, node_id, 2.0) === 0.5f0
+    @test low_storage_factor([1.0], edge_metadata, node_id, 2.0) === 0.5
+    @test low_storage_factor([3.0f0], edge_metadata, node_id, 2.0) === 1.0f0
+    @test low_storage_factor([3.0], edge_metadata, node_id, 2.0) === 1.0
 end
 
 @testitem "constraints_from_nodes" begin
diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index 610407f6c..c9ad2b9d4 100644
--- a/core/test/validation_test.jl
+++ b/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, 0, edge_type, 0, (id_1, id_2), (0, 0))
+        graph[id_1, id_2] = EdgeMetadata(0, 0, edge_type, 0, (id_1, id_2), -1, -1)
         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, 0, edge_type, 0, (id_1, id_2), (0, 0))
+        graph[id_1, id_2] = EdgeMetadata(0, 0, edge_type, 0, (id_1, id_2), -1, -1)
         return nothing
     end
 
@@ -462,12 +462,12 @@ end
     model = Ribasim.Model(toml_path)
     (; graph, basin) = model.integrator.p
     for edge_metadata in values(graph.edge_data)
-        (; edge, basin_idxs) = edge_metadata
+        (; edge, basin_idx_src, basin_idx_dst) = edge_metadata
         id_src, id_dst = edge
         if id_src.type == NodeType.Basin
-            @test id_src == basin.node_id.values[basin_idxs[1]]
+            @test id_src == basin.node_id.values[basin_idx_src]
         elseif id_dst.type == NodeType.Basin
-            @test id_dst == basin.node_id.values[basin_idxs[2]]
+            @test id_dst == basin.node_id.values[basin_idx_dst]
         end
     end
 end

From 4e6a085282746c7eb4265dbe02c566e4661498f7 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 13 May 2024 15:04:19 +0200
Subject: [PATCH 15/15] Revert https://github.com/Deltares/Ribasim/pull/1458,
 causes CI to fail (but not locally?)

---
 core/Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/Project.toml b/core/Project.toml
index 45441c4ed..1a9fb12e3 100644
--- a/core/Project.toml
+++ b/core/Project.toml
@@ -51,7 +51,7 @@ ComponentArrays = "0.13,0.14,0.15"
 Configurations = "0.17"
 DBInterface = "2.4"
 DataFrames = "1.4"
-DataInterpolations = "4.4, 5"
+DataInterpolations = "4.4"
 DataStructures = "0.18"
 Dates = "<0.0.1,1"
 Dictionaries = "0.3.25, 0.4"