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 a29c7d82f..8992b9a9a 100644 --- a/core/src/solve.jl +++ b/core/src/solve.jl @@ -975,7 +975,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