Move setting previous flows before timestep to SciMLBase.step

core/src/bmi.jl

@@ -18,13 +18,7 @@ function BMI.update(model::Model)::Model
 end
 
 function BMI.update_until(model::Model, time::Float64)::Model
-    # Recompute flow_prev to take into account flow discontinuities introduced
-    # by parameter changes through BMI
-    (; integrator) = model
-    (; u, p, t) = integrator
-    (; flow, flow_prev) = p.graph[]
-    water_balance!(get_du(integrator), u, p, t)
-    copyto!(flow_prev, get_tmp(flow, 0))
+    (; t) = model.integrator
     dt = time - t
     if dt < 0
         error("The model has already passed the given timestamp.")

core/src/callback.jl

@@ -139,7 +139,6 @@ function integrate_flows!(u, t, integrator)::Nothing
                 0.5 * (get_flow(graph, edge..., 0) + get_flow_prev(graph, edge..., 0)) * dt
         end
     end
-
     copyto!(flow_prev, flow)
     copyto!(vertical_flux_prev, vertical_flux)
     return nothing

core/src/model.jl

@@ -213,6 +213,7 @@ function SciMLBase.step!(model::Model, dt::Float64)::Model
     if round(ntimes) ≈ ntimes
         update_allocation!(integrator)
     end
+    set_previous_flows!(integrator)
     step!(integrator, dt, true)
     return model
 end

core/src/util.jl

@@ -749,3 +749,20 @@ end
 function basin_areas(basin::Basin, state_idx::Int)
     return basin.level_to_area[state_idx].u
 end
+
+"""
+Just before setting a timestep, call water_balance! again
+to get a correct value of the flows for integrating
+"""
+function set_previous_flows!(integrator)
+    (; p, u, t) = integrator
+    (; flow, flow_prev) = p.graph[]
+    (; vertical_flux, vertical_flux_prev) = p.basin
+    du = get_du(integrator)
+    water_balance!(du, u, p, t)
+
+    flow = get_tmp(flow, 0)
+    vertical_flux = get_tmp(vertical_flux, 0)
+    copyto!(flow_prev, flow)
+    copyto!(vertical_flux_prev, vertical_flux)
+end

core/test/bmi_test.jl

@@ -87,6 +87,7 @@ end
     model = Ribasim.Model(config)
     demand = BMI.get_value_ptr(model, "user_demand.demand")
     realized = BMI.get_value_ptr(model, "user_demand.realized")
+    # One year in seconds
     year = model.integrator.p.user_demand.demand_itp[2][1].t[2]
     demand_start = 1e-3
     slope = 1e-3 / year