From 415158f24c59cff7da4006476737a3d876d0b620 Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Fri, 8 Sep 2023 20:54:43 +0200
Subject: [PATCH] avoid allocating value vectors in get_level (#582)

`Dictionary` uses `Indices{I}` as keys, and `Vector{T}` as values. The
Parameters contain both, and therefore it was free to construct a
`Dictionary` in a frequently called function like `get_level`. However
with autodiff, the values could be a ReinterpretArray with Duals instead
of just a Vector. This meant that on Dictionary creation it would
convert the ReinterpretArray to a Vector, leading to many allocations.

This is on top of https://github.com/Deltares/Ribasim/pull/581. After
that, this was responsible for 94% of the time spent. With this PR that
goes down to about 2%, leading to a nice little speedup.

---------

Co-authored-by: Hofer-Julian <30049909+Hofer-Julian@users.noreply.github.com>
---
 core/src/Ribasim.jl |  2 +-
 core/src/solve.jl   |  2 +-
 core/src/utils.jl   | 32 +++++++++++++++++---------------
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl
index 78b615aaa..8821ef2c8 100644
--- a/core/src/Ribasim.jl
+++ b/core/src/Ribasim.jl
@@ -12,7 +12,7 @@ using ComponentArrays: ComponentVector
 using DataInterpolations: LinearInterpolation, derivative
 using Dates
 using DBInterface: execute, prepare
-using Dictionaries: Indices, Dictionary, gettoken, gettokenvalue, dictionary
+using Dictionaries: Indices, Dictionary, gettoken, dictionary
 using ForwardDiff: pickchunksize
 using DiffEqCallbacks
 using Graphs: DiGraph, add_edge!, adjacency_matrix, inneighbors, outneighbors
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 30b4297ce..d7b9c65f3 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -1037,7 +1037,7 @@ end
 function formulate_flows!(
     p::Parameters,
     storage::AbstractVector,
-    current_level,
+    current_level::AbstractVector,
     flow::AbstractMatrix,
     t::Float64,
 )::Nothing
diff --git a/core/src/utils.jl b/core/src/utils.jl
index 13da0e4b6..d4b80baaa 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -406,34 +406,36 @@ end
 Get the current water level of a node ID.
 The ID can belong to either a Basin or a LevelBoundary.
 """
-function get_level(p::Parameters, node_id::Int, current_level, t::Float64)::Real
+function get_level(
+    p::Parameters,
+    node_id::Int,
+    current_level::AbstractVector,
+    t::Float64,
+)::Real
     (; basin, level_boundary) = p
-    # since the node_id fields are already Indices, Dictionary creation is instant
-    basin = Dictionary(basin.node_id, current_level)
-    hasindex, token = gettoken(basin, node_id)
+    hasindex, i = id_index(basin.node_id, node_id)
     return if hasindex
-        gettokenvalue(basin, token)
+        current_level[i]
     else
-        boundary = Dictionary(level_boundary.node_id, level_boundary.level)
-        boundary[node_id](t)
+        i = findsorted(level_boundary.node_id, node_id)::Int
+        level_boundary.level[i](t)
     end
 end
 
 "Get the index of an ID in a set of indices."
-function id_index(ids::Indices{Int}, id::Int)
-    # There might be a better approach for this, this feels too internal
-    # the second return is the token, a Tuple{Int, Int}
-    hasindex, (_, idx) = gettoken(ids, id)
-    return hasindex, idx
+function id_index(ids::Indices{Int}, id::Int)::Tuple{Bool, Int}
+    # We avoid creating Dictionary here since it converts the values to a Vector,
+    # leading to allocations when used with PreallocationTools's ReinterpretArrays.
+    hasindex, (_, i) = gettoken(ids, id)
+    return hasindex, i
 end
 
 "Return the bottom elevation of the basin with index i, or nothing if it doesn't exist"
 function basin_bottom(basin::Basin, node_id::Int)::Union{Float64, Nothing}
-    basin = Dictionary(basin.node_id, basin.level)
-    hasindex, token = gettoken(basin, node_id)
+    hasindex, i = id_index(basin.node_id, node_id)
     return if hasindex
         # get level(storage) interpolation function
-        level_discrete = gettokenvalue(basin, token)
+        level_discrete = basin.level[i]
         # and return the first level in this vector, representing the bottom
         first(level_discrete)
     else