Finetune step_limiter!

core/src/allocation_optim.jl

@@ -311,7 +311,7 @@ function get_basin_data(
     (; mean_input_flows) = allocation
     @assert node_id.type == NodeType.Basin
     influx = mean_input_flows[(node_id, node_id)][]
-    storage_basin = basin.current_storage[parent(u)][node_id.idx]
+    storage_basin = basin.current_properties.current_storage[parent(u)][node_id.idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)
     if isempty(control_inneighbors)
         level_demand_idx = 0

core/src/bmi.jl

@@ -33,9 +33,9 @@ end
 function BMI.get_value_ptr(model::Model, name::AbstractString)::AbstractVector{Float64}
     (; u, p) = model.integrator
     if name == "basin.storage"
-        p.basin.current_storage[parent(u)]
+        p.basin.current_properties.current_storage[parent(u)]
     elseif name == "basin.level"
-        p.basin.current_level[parent(u)]
+        p.basin.current_properties.current_level[parent(u)]
     elseif name == "basin.infiltration"
         p.basin.vertical_flux.infiltration
     elseif name == "basin.drainage"

core/src/callback.jl

@@ -316,8 +316,10 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
     end
 
     # Update the Basin concentrations again based on the removed mass
-    basin.concentration_state .= basin.mass ./ basin.current_storage[parent(u)]
-    basin.storage_prev .= basin.current_storage[parent(u)]
+    basin.concentration_state .=
+        basin.mass ./ basin.current_properties.current_storage[parent(u)]
+    basin.storage_prev .= basin.current_properties.current_storage[parent(u)]
+    basin.level_prev .= basin.current_properties.current_level[parent(u)]
 
     return nothing
 end
@@ -361,8 +363,8 @@ function save_basin_state(u, t, integrator)
     (; p) = integrator
     (; basin) = p
     du = get_du(integrator)
-    current_storage = basin.current_storage[parent(du)]
-    current_level = basin.current_level[parent(du)]
+    current_storage = basin.current_properties.current_storage[parent(du)]
+    current_level = basin.current_properties.current_level[parent(du)]
     water_balance!(du, u, p, t)
     SavedBasinState(; storage = copy(current_storage), level = copy(current_level), t)
 end
@@ -448,7 +450,7 @@ function check_water_balance_error!(
     (; u, p, t) = integrator
     (; basin, water_balance_abstol, water_balance_reltol) = p
     errors = false
-    current_storage = basin.current_storage[parent(u)]
+    current_storage = basin.current_properties.current_storage[parent(u)]
 
     # The initial storage is irrelevant for the storage rate and can only cause
     # floating point truncation errors
@@ -514,7 +516,8 @@ end
 
 function check_negative_storage(u, t, integrator)::Nothing
     (; basin) = integrator.p
-    (; node_id, current_storage) = basin
+    (; node_id, current_properties) = basin
+    (; current_storage) = current_properties
     du = get_du(integrator)
     set_current_basin_properties!(du, u, integrator.p, t)
     current_storage = current_storage[parent(du)]
@@ -737,7 +740,7 @@ end
 function update_subgrid_level!(integrator)::Nothing
     (; p) = integrator
     du = get_du(integrator)
-    basin_level = p.basin.current_level[parent(du)]
+    basin_level = p.basin.current_properties.current_level[parent(du)]
     subgrid = integrator.p.subgrid
     for (i, (index, interp)) in enumerate(zip(subgrid.basin_index, subgrid.interpolations))
         subgrid.level[i] = interp(basin_level[index])
@@ -823,7 +826,7 @@ update_userd_conc!(integrator)::Nothing =
 function update_allocation!(integrator)::Nothing
     (; p, t, u) = integrator
     (; allocation, basin) = p
-    (; current_storage) = basin
+    (; current_storage) = basin.current_properties
     (; allocation_models, mean_input_flows, mean_realized_flows) = allocation
 
     # Make sure current storages are up to date

core/src/model.jl

@@ -118,6 +118,8 @@ function Model(config::Config)::Model
 
     # Synchronize level with storage
     set_current_basin_properties!(du0, u0, parameters, t0)
+    parameters.basin.level_prev .=
+        parameters.basin.current_properties.current_level[parent(u0)]
 
     saveat = convert_saveat(config.solver.saveat, t_end)
     saveat isa Float64 && push!(tstops, range(0, t_end; step = saveat))

core/src/parameter.jl

@@ -1,3 +1,6 @@
+const LOW_STORAGE_THRESHOLD = 10.0
+const USER_DEMAND_MIN_LEVEL_THRESHOLD = 0.1
+
 const SolverStats = @NamedTuple{
     time::Float64,
     rhs_calls::Int,
@@ -300,6 +303,21 @@ In-memory storage of saved instantaneous storages and levels for writing to resu
     t::Float64
 end
 
+"""
+Caches of current basin properties
+"""
+struct CurrentBasinProperties
+    current_storage::Cache
+    current_low_storage_factor::Cache
+    current_level::Cache
+    current_area::Cache
+    current_cumulative_precipitation::Cache
+    current_cumulative_drainage::Cache
+    function CurrentBasinProperties(n)
+        new((cache(n) for _ in 1:6)...)
+    end
+end
+
 """
 Requirements:
 
@@ -332,11 +350,7 @@ end
     cumulative_precipitation_saveat::Vector{Float64} = zeros(length(node_id))
     cumulative_drainage_saveat::Vector{Float64} = zeros(length(node_id))
     # Cache this to avoid recomputation
-    current_storage::Cache = cache(length(node_id))
-    current_level::Cache = cache(length(node_id))
-    current_area::Cache = cache(length(node_id))
-    current_cumulative_precipitation::Cache = cache(length(node_id))
-    current_cumulative_drainage::Cache = cache(length(node_id))
+    current_properties::CurrentBasinProperties = CurrentBasinProperties(length(node_id))
     # Discrete values for interpolation
     storage_to_level::Vector{
         LinearInterpolationIntInv{
@@ -355,6 +369,8 @@ end
     # Data source for concentration updates
     concentration_time::StructVector{BasinConcentrationV1, D, Int}
 
+    # Level for each Basin at the previous time step
+    level_prev::Vector{Float64} = zeros(length(node_id))
     # Concentrations
     # Config setting to enable/disable evaporation of mass
     evaporate_mass::Bool = true

core/src/read.jl

@@ -574,13 +574,10 @@ end
 function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
     node_id = get_ids(db, "Basin")
     n = length(node_id)
-    current_level = cache(n)
-    current_area = cache(n)
 
     evaporate_mass = config.solver.evaporate_mass
     precipitation = zeros(n)
     potential_evaporation = zeros(n)
-    evaporation = zeros(n)
     drainage = zeros(n)
     infiltration = zeros(n)
     table = (; precipitation, potential_evaporation, drainage, infiltration)
@@ -686,8 +683,6 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
         inflow_ids = [collect(inflow_ids(graph, id)) for id in node_id],
         outflow_ids = [collect(outflow_ids(graph, id)) for id in node_id],
         vertical_flux,
-        current_level,
-        current_area,
         storage_to_level,
         level_to_area,
         demand,