From d8fe0a1b98f4a1d2f9fc97f9e0d9fcab052f0fcf Mon Sep 17 00:00:00 2001
From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com>
Date: Tue, 27 Aug 2024 11:59:39 +0200
Subject: [PATCH] refactor `get_influx` for performance (#1753)

This is a follow-up of https://github.com/Deltares/Ribasim/pull/1606.

Hopefully fixes the runtime of the HWS model.
---
 core/src/util.jl | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/core/src/util.jl b/core/src/util.jl
index bccc037e6..42bd136ea 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -540,12 +540,21 @@ function get_influx(basin::Basin, node_id::NodeID)::Float64
 end
 
 function get_influx(basin::Basin, basin_idx::Int; prev::Bool = false)::Float64
-    (; vertical_flux, vertical_flux_prev, vertical_flux_from_input) = basin
-    vertical_flux = wrap_forcing(vertical_flux[parent(vertical_flux_from_input)])
-    flux_vector = prev ? vertical_flux_prev : vertical_flux
-    (; precipitation, evaporation, drainage, infiltration) = flux_vector
-    return precipitation[basin_idx] - evaporation[basin_idx] + drainage[basin_idx] -
-           infiltration[basin_idx]
+    (; node_id, vertical_flux, vertical_flux_prev, vertical_flux_from_input) = basin
+    influx = if prev
+        vertical_flux_prev.precipitation[basin_idx] -
+        vertical_flux_prev.evaporation[basin_idx] +
+        vertical_flux_prev.drainage[basin_idx] -
+        vertical_flux_prev.infiltration[basin_idx]
+    else
+        n = length(node_id)
+        vertical_flux = vertical_flux[parent(vertical_flux_from_input)]
+        vertical_flux[basin_idx] - # precipitation
+        vertical_flux[n + basin_idx] + # evaporation
+        vertical_flux[2n + basin_idx] - # drainage
+        vertical_flux[3n + basin_idx] # infiltration
+    end
+    return influx
 end
 
 inflow_edge(graph, node_id)::EdgeMetadata = graph[inflow_id(graph, node_id), node_id]