Precalculate neighboring edges for UserDemand

Deltares · May 8, 2024 · ad544bb · ad544bb
1 parent d2fef95
commit ad544bb
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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