Save all results via SavingCallback (#1829)

Fixes #1825.

core/src/callback.jl

@@ -197,7 +197,7 @@ function save_basin_state(u, t, integrator)
     current_storage = basin.current_storage[parent(du)]
     current_level = basin.current_level[parent(du)]
     water_balance!(du, u, p, t)
-    SavedBasinState(; storage = copy(current_storage), level = copy(current_level))
+    SavedBasinState(; storage = copy(current_storage), level = copy(current_level), t)
 end
 
 """
@@ -206,10 +206,13 @@ Both computed by the solver and integrated exactly. Also computes the total hori
 inflow and outflow per basin.
 """
 function save_flow(u, t, integrator)
-    (; p, sol) = integrator
-    (; basin, state_inflow_edge, state_outflow_edge, flow_boundary) = p
+    (; p) = integrator
+    (; basin, state_inflow_edge, state_outflow_edge, flow_boundary, u_prev_saveat) = p
     Δt = get_Δt(integrator)
-    flow_mean = (u - sol(t - Δt)) / Δt
+    flow_mean = (u - u_prev_saveat) / Δt
+
+    # Current u is previous u in next computation
+    u_prev_saveat .= u
 
     inflow_mean = zeros(length(basin.node_id))
     outflow_mean = zeros(length(basin.node_id))
@@ -262,6 +265,7 @@ function save_flow(u, t, integrator)
         flow_boundary = flow_boundary_mean,
         precipitation,
         drainage,
+        t,
     )
     check_water_balance_error(integrator, saved_flow, Δt)
     return saved_flow

core/src/model.jl

@@ -104,6 +104,7 @@ function Model(config::Config)::Model
     du0 = zero(u0)
 
     parameters = set_state_flow_edges(parameters, u0)
+    parameters = @set parameters.u_prev_saveat = zero(u0)
 
     # The Solver algorithm
     alg = algorithm(config.solver; u0)
@@ -150,10 +151,10 @@ function Model(config::Config)::Model
         progress = true,
         progress_name = "Simulating",
         progress_steps = 100,
+        save_everystep = false,
         callback,
         tstops,
         isoutofdomain,
-        saveat,
         adaptive,
         dt,
         config.solver.dtmin,
@@ -175,7 +176,8 @@ function Model(config::Config)::Model
 end
 
 "Get all saved times in seconds since start"
-tsaves(model::Model)::Vector{Float64} = model.integrator.sol.t
+tsaves(model::Model)::Vector{Float64} =
+    [0.0, (cvec.t for cvec in model.saved.flow.saveval)...]
 
 "Get all saved times as a Vector{DateTime}"
 function datetimes(model::Model)::Vector{DateTime}

core/src/parameter.jl

@@ -268,6 +268,7 @@ In-memory storage of saved mean flows for writing to results.
 - `flow_boundary`: The exact integrated mean flows of flow boundaries
 - `precipitation`: The exact integrated mean precipitation
 - `drainage`: The exact integrated mean drainage
+- `t`: Endtime of the interval over which is averaged
 """
 @kwdef struct SavedFlow{V}
     flow::V
@@ -276,6 +277,7 @@ In-memory storage of saved mean flows for writing to results.
     flow_boundary::Vector{Float64}
     precipitation::Vector{Float64}
     drainage::Vector{Float64}
+    t::Float64
 end
 
 """
@@ -284,6 +286,7 @@ In-memory storage of saved instantaneous storages and levels for writing to resu
 @kwdef struct SavedBasinState
     storage::Vector{Float64}
     level::Vector{Float64}
+    t::Float64
 end
 
 """
@@ -815,7 +818,7 @@ const ModelGraph = MetaGraph{
     Float64,
 }
 
-@kwdef struct Parameters{C1, C2, C3, C4, V1, V2}
+@kwdef struct Parameters{C1, C2, C3, C4, C5, V1, V2}
     starttime::DateTime
     graph::ModelGraph
     allocation::Allocation
@@ -843,6 +846,8 @@ const ModelGraph = MetaGraph{
     # Water balance tolerances
     water_balance_abstol::Float64
     water_balance_reltol::Float64
+    # State at previous saveat
+    u_prev_saveat::C5 = ComponentVector()
 end
 
 # To opt-out of type checking for ForwardDiff

core/test/docs.toml

@@ -33,8 +33,8 @@ dtmax = 0.0         # optional, default length of simulation
 force_dtmin = false # optional, default false
 abstol = 1e-6       # optional, default 1e-6
 reltol = 1e-5       # optional, default 1e-5
-water_balance_abstol = 1e-6 # optional, default 1e-6
-water_balance_reltol = 1e-6 # optional, default 1e-6
+water_balance_abstol = 1e-3 # optional, default 1e-3
+water_balance_reltol = 1e-2 # optional, default 1e-2
 maxiters = 1e9      # optional, default 1e9
 sparse = true       # optional, default true
 autodiff = true     # optional, default true

core/test/run_models_test.jl

@@ -203,7 +203,7 @@ end
     @test all(isconcretetype, fieldtypes(typeof(p)))
 
     @test successful_retcode(model)
-    @test allunique(Ribasim.tsaves(model))
+    @test length(model.integrator.sol) == 2 # start and end
     @test model.integrator.p.basin.current_storage[Float64[]] ≈
           Float32[803.7093, 803.68274, 495.241, 1318.3053] skip = Sys.isapple() atol = 1.5