From b8692cd3f7bd156448a3ba8463b22431843d1f0e Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 21 Feb 2024 13:52:38 +0100
Subject: [PATCH 01/19] Mean output flows

---
 core/src/callback.jl  | 43 ++++++++++++++++++++++++++++++++++++++-----
 core/src/graph.jl     | 10 ++++++++++
 core/src/parameter.jl |  5 +++++
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index e5d8793a7..879d8d5d6 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -7,7 +7,7 @@ function set_initial_discrete_controlled_parameters!(
     storage0::Vector{Float64},
 )::Nothing
     (; p) = integrator
-    (; basin, discrete_control) = p
+    (; discrete_control) = p
 
     n_conditions = length(discrete_control.condition_value)
     condition_diffs = zeros(Float64, n_conditions)
@@ -40,6 +40,9 @@ function create_callbacks(
     basin_cb = PresetTimeCallback(tstops, update_basin)
     push!(callbacks, basin_cb)
 
+    integrating_flows_cb = FunctionCallingCallback(integrate_flows!)
+    push!(callbacks, integrating_flows_cb)
+
     tstops = get_tstops(tabulated_rating_curve.time.time, starttime)
     tabulated_rating_curve_cb = PresetTimeCallback(tstops, update_tabulated_rating_curve!)
     push!(callbacks, tabulated_rating_curve_cb)
@@ -87,6 +90,29 @@ function create_callbacks(
     return callback, saved
 end
 
+function integrate_flows!(u, t, integrator)::Nothing
+    (; p, tprev) = integrator
+    (; graph) = p
+    (;
+        flow,
+        flow_vertical,
+        flow_prev,
+        flow_vertical_prev,
+        flow_integrated,
+        flow_vertical_integrated,
+    ) = graph[]
+    flow = get_tmp(flow, 0)
+    flow_vertical = get_tmp(flow_vertical, 0)
+    Δt = t - tprev
+
+    @. flow_integrated += 0.5 * (flow + flow_prev) * Δt
+    @. flow_vertical_integrated += 0.5 * (flow_vertical + flow_vertical_prev) * Δt
+
+    copyto!(flow_prev, flow)
+    copyto!(flow_vertical_prev, flow_vertical)
+    return nothing
+end
+
 """
 Listens for changes in condition truths.
 """
@@ -379,10 +405,17 @@ end
 
 "Copy the current flow to the SavedValues"
 function save_flow(u, t, integrator)
-    vcat(
-        get_tmp(integrator.p.graph[].flow_vertical, 0.0),
-        get_tmp(integrator.p.graph[].flow, 0.0),
-    )
+    (; graph) = integrator.p
+    (; flow_integrated, flow_vertical_integrated, tprev_flow_save) = graph[]
+    Δt = t - tprev_flow_save[1]
+    tprev_flow_save[1] = t
+
+    mean_flow_vertical = flow_vertical_integrated / Δt
+    mean_flow = flow_integrated / Δt
+
+    fill!(flow_vertical_integrated, 0.0)
+    fill!(flow_integrated, 0.0)
+    return vcat(mean_flow_vertical, mean_flow)
 end
 
 function update_subgrid_level!(integrator)::Nothing
diff --git a/core/src/graph.jl b/core/src/graph.jl
index a7b69eaf8..e901fbab7 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -90,19 +90,29 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
     end
 
     flow = zeros(flow_counter)
+    flow_prev = zeros(flow_counter)
+    flow_integrated = zeros(flow_counter)
     flow_vertical = zeros(flow_vertical_counter)
+    flow_vertical_prev = zeros(flow_vertical_counter)
+    flow_vertical_integrated = zeros(flow_vertical_counter)
     if config.solver.autodiff
         flow = DiffCache(flow, chunk_sizes)
         flow_vertical = DiffCache(flow_vertical, chunk_sizes)
     end
+    tprev_flow_save = zeros(1)
     graph_data = (;
         node_ids,
         edge_ids,
         edges_source,
         flow_dict,
         flow,
+        flow_prev,
+        flow_integrated,
         flow_vertical_dict,
         flow_vertical,
+        flow_vertical_prev,
+        flow_vertical_integrated,
+        tprev_flow_save,
     )
     graph = @set graph.graph_data = graph_data
 
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index fc8c8feb7..da5dd8de9 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -553,8 +553,13 @@ struct Parameters{T, C1, C2}
             edges_source::Dict{Int, Set{EdgeMetadata}},
             flow_dict::Dict{Tuple{NodeID, NodeID}, Int},
             flow::T,
+            flow_prev::Vector{Float64},
+            flow_integrated::Vector{Float64},
             flow_vertical_dict::Dict{NodeID, Int},
             flow_vertical::T,
+            flow_vertical_prev::Vector{Float64},
+            flow_vertical_integrated::Vector{Float64},
+            tprev_flow_save::Vector{Float64},
         },
         MetaGraphsNext.var"#11#13",
         Float64,

From f9d0eba64744f015413df77868deae5cff008010 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 21 Feb 2024 18:38:58 +0100
Subject: [PATCH 02/19] Fix test and add test

---
 core/test/run_models_test.jl | 13 +++++++++++++
 core/test/time_test.jl       |  2 +-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 051a27d75..110272fcb 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -476,3 +476,16 @@ end
         0,
     ) ≈ 5.0 atol = 0.001 skip = Sys.isapple()
 end
+
+@testitem "mean_flow" begin
+    toml_path =
+        normpath(@__DIR__, "../../generated_testmodels/linear_resistance/ribasim.toml")
+    @test ispath(toml_path)
+    t_end = 3.16224e7
+
+    config = Ribasim.Config(toml_path; solver_saveat = t_end)
+    model = Ribasim.run(config)
+    mean_flow = model.saved.flow.saveval[1]
+    storage = Ribasim.get_storages_and_levels(model).storage[1, :]
+    @test storage[2] ≈ storage[1] - t_end * mean_flow[3] atol = 1
+end
diff --git a/core/test/time_test.jl b/core/test/time_test.jl
index 82ad216d3..5ff0ed97a 100644
--- a/core/test/time_test.jl
+++ b/core/test/time_test.jl
@@ -29,5 +29,5 @@
     t = Ribasim.seconds_since.(flow_1_to_2.time, model.config.starttime)
     flow_expected = @. 1 + sin(0.5 * π * (t - t[1]) / (t[end] - t[1]))^2
     # some difference is expected since the modeled flow is for the period up to t
-    @test isapprox(flow_added_1, flow_expected, rtol = 0.005)
+    @test isapprox(flow_added_1, flow_expected, rtol = 0.01)
 end

From 83ec27870e7002dae753d478a40a7d5500ce6080 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 22 Feb 2024 15:43:51 +0100
Subject: [PATCH 03/19] =?UTF-8?q?Update=20flow=20save=20=CE=94t?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 core/src/callback.jl | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 879d8d5d6..d2bde6e10 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -40,7 +40,7 @@ function create_callbacks(
     basin_cb = PresetTimeCallback(tstops, update_basin)
     push!(callbacks, basin_cb)
 
-    integrating_flows_cb = FunctionCallingCallback(integrate_flows!)
+    integrating_flows_cb = FunctionCallingCallback(integrate_flows!; func_start = false)
     push!(callbacks, integrating_flows_cb)
 
     tstops = get_tstops(tabulated_rating_curve.time.time, starttime)
@@ -405,10 +405,19 @@ end
 
 "Copy the current flow to the SavedValues"
 function save_flow(u, t, integrator)
-    (; graph) = integrator.p
-    (; flow_integrated, flow_vertical_integrated, tprev_flow_save) = graph[]
-    Δt = t - tprev_flow_save[1]
-    tprev_flow_save[1] = t
+    (; tprev, p) = integrator
+    (; graph) = p
+    (; flow_integrated, flow_vertical_integrated, saveat) = graph[]
+
+    Δt = if iszero(saveat)
+        t - tprev
+    elseif isinf(saveat)
+        t
+    else
+        t_end = integrator.sol.prob.tspan[2]
+        # The last interval might be shorter than saveat
+        t_end - t >= saveat ? saveat : t - tprev
+    end
 
     mean_flow_vertical = flow_vertical_integrated / Δt
     mean_flow = flow_integrated / Δt

From 9257ffbe2dbfb00737609a0f81a50280ead87699 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 22 Feb 2024 15:45:45 +0100
Subject: [PATCH 04/19] Add saveat to tstops

---
 core/src/graph.jl | 3 +--
 core/src/model.jl | 8 ++++++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/core/src/graph.jl b/core/src/graph.jl
index e901fbab7..9cfda9a80 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -99,7 +99,6 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         flow = DiffCache(flow, chunk_sizes)
         flow_vertical = DiffCache(flow_vertical, chunk_sizes)
     end
-    tprev_flow_save = zeros(1)
     graph_data = (;
         node_ids,
         edge_ids,
@@ -112,7 +111,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         flow_vertical,
         flow_vertical_prev,
         flow_vertical_integrated,
-        tprev_flow_save,
+        config.solver.saveat,
     )
     graph = @set graph.graph_data = graph_data
 
diff --git a/core/src/model.jl b/core/src/model.jl
index c6e6bc787..9ad1f0730 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -46,7 +46,7 @@ function Model(config::Config)::Model
     # All data from the database that we need during runtime is copied into memory,
     # so we can directly close it again.
     db = SQLite.DB(db_path)
-    local parameters, state, n, tstops
+    local parameters, state, n, tstops, tstops_flow_boundary, tstops_user
     try
         parameters = Parameters(db, config)
 
@@ -79,7 +79,6 @@ function Model(config::Config)::Model
         tstops_flow_boundary = get_tstops(time_flow_boundary.time, config.starttime)
         time_user = load_structvector(db, config, UserTimeV1)
         tstops_user = get_tstops(time_user.time, config.starttime)
-        tstops = sort(unique(vcat(tstops_flow_boundary, tstops_user)))
 
         # use state
         state = load_structvector(db, config, BasinStateV1)
@@ -107,6 +106,11 @@ function Model(config::Config)::Model
     timespan = (t0, t_end)
     saveat = convert_saveat(config.solver.saveat, t_end)
 
+    tstops_saveat =
+        isfinite(config.solver.saveat) ? range(t0, t_end; step = config.solver.saveat) :
+        Float64[]
+    tstops = sort(unique(vcat(tstops_flow_boundary, tstops_user, tstops_saveat)))
+
     jac_prototype = config.solver.sparse ? get_jac_prototype(parameters) : nothing
     RHS = ODEFunction(water_balance!; jac_prototype)
 

From edfe8cdd9c92db18ac623c1ccfa8f8b14a9ea520 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 22 Feb 2024 15:48:08 +0100
Subject: [PATCH 05/19] small stuff

---
 core/src/parameter.jl  |  2 +-
 core/src/solve.jl      | 12 ------------
 core/test/time_test.jl |  2 +-
 3 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index eecb14c60..873240ac1 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -564,7 +564,7 @@ struct Parameters{T, C1, C2}
             flow_vertical::T,
             flow_vertical_prev::Vector{Float64},
             flow_vertical_integrated::Vector{Float64},
-            tprev_flow_save::Vector{Float64},
+            saveat::Float64,
         },
         MetaGraphsNext.var"#11#13",
         Float64,
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 4a8a04aa6..5c662952b 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -665,15 +665,3 @@ function formulate_flows!(p::Parameters, storage::AbstractVector, t::Number)::No
     formulate_flow!(level_boundary, p, storage, t)
     formulate_flow!(terminal, p, storage, t)
 end
-
-function track_waterbalance!(u, t, integrator)::Nothing
-    (; p, tprev, uprev) = integrator
-    dt = t - tprev
-    du = u - uprev
-    p.storage_diff .+= du
-    p.precipitation.total .+= p.precipitation.value .* dt
-    p.evaporation.total .+= p.evaporation.value .* dt
-    p.infiltration.total .+= p.infiltration.value .* dt
-    p.drainage.total .+= p.drainage.value .* dt
-    return nothing
-end
diff --git a/core/test/time_test.jl b/core/test/time_test.jl
index 5ff0ed97a..82ad216d3 100644
--- a/core/test/time_test.jl
+++ b/core/test/time_test.jl
@@ -29,5 +29,5 @@
     t = Ribasim.seconds_since.(flow_1_to_2.time, model.config.starttime)
     flow_expected = @. 1 + sin(0.5 * π * (t - t[1]) / (t[end] - t[1]))^2
     # some difference is expected since the modeled flow is for the period up to t
-    @test isapprox(flow_added_1, flow_expected, rtol = 0.01)
+    @test isapprox(flow_added_1, flow_expected, rtol = 0.005)
 end

From 64bbac061a07a76bd9ca0e50c16dbade6c823e7f Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 22 Feb 2024 16:54:56 +0100
Subject: [PATCH 06/19] Add tests

---
 core/src/callback.jl         | 25 +++++++++++++++++++++----
 core/src/config.jl           | 18 +++++++++++++++---
 core/src/model.jl            | 12 +++++-------
 core/test/run_models_test.jl | 25 +++++++++++++++++--------
 4 files changed, 58 insertions(+), 22 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index d2bde6e10..18edbcf97 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -31,7 +31,7 @@ Returns the CallbackSet and the SavedValues for flow.
 function create_callbacks(
     parameters::Parameters,
     config::Config;
-    saveat,
+    saveat_flow,
 )::Tuple{CallbackSet, SavedResults}
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
     callbacks = SciMLBase.DECallback[]
@@ -58,7 +58,8 @@ function create_callbacks(
 
     # save the flows over time, as a Vector of the nonzeros(flow)
     saved_flow = SavedValues(Float64, Vector{Float64})
-    save_flow_cb = SavingCallback(save_flow, saved_flow; saveat, save_start = false)
+    save_flow_cb =
+        SavingCallback(save_flow, saved_flow; saveat = saveat_flow, save_start = false)
     push!(callbacks, save_flow_cb)
 
     # interpolate the levels
@@ -104,9 +105,23 @@ function integrate_flows!(u, t, integrator)::Nothing
     flow = get_tmp(flow, 0)
     flow_vertical = get_tmp(flow_vertical, 0)
     Δt = t - tprev
+    @show Δt
 
-    @. flow_integrated += 0.5 * (flow + flow_prev) * Δt
-    @. flow_vertical_integrated += 0.5 * (flow_vertical + flow_vertical_prev) * Δt
+    flow_effective = if length(integrator.sol.t) == 2
+        # If flow_prev is not populated yet
+        flow
+    else
+        0.5 * (flow + flow_prev)
+    end
+
+    flow_vertical_effective = if length(integrator.sol.t) == 2
+        flow_vertical
+    else
+        0.5 * (flow_vertical + flow_vertical_prev)
+    end
+
+    @. flow_integrated += flow_effective * Δt
+    @. flow_vertical_integrated += flow_vertical_effective * Δt
 
     copyto!(flow_prev, flow)
     copyto!(flow_vertical_prev, flow_vertical)
@@ -419,6 +434,8 @@ function save_flow(u, t, integrator)
         t_end - t >= saveat ? saveat : t - tprev
     end
 
+    @show Δt
+
     mean_flow_vertical = flow_vertical_integrated / Δt
     mean_flow = flow_integrated / Δt
 
diff --git a/core/src/config.jl b/core/src/config.jl
index a8aabb573..f75845230 100644
--- a/core/src/config.jl
+++ b/core/src/config.jl
@@ -232,16 +232,28 @@ function algorithm(solver::Solver)::OrdinaryDiffEqAlgorithm
 end
 
 "Convert the saveat Float64 from our Config to SciML's saveat"
-function convert_saveat(saveat::Float64, t_end::Float64)::Union{Float64, Vector{Float64}}
+function convert_saveat(
+    saveat::Float64,
+    t_end::Float64;
+    state::Bool = true,
+)::Union{Float64, Vector{Float64}}
     if iszero(saveat)
         # every step
         Float64[]
     elseif saveat == Inf
         # only the start and end
-        [0.0, t_end]
+        if state
+            [0.0, t_end]
+        else
+            [t_end]
+        end
     elseif isfinite(saveat)
         # every saveat seconds
-        saveat
+        if state
+            saveat
+        else
+            collect(range(0.0, t_end; step = saveat))
+        end
     else
         @error "Invalid saveat" saveat
         error("Invalid saveat")
diff --git a/core/src/model.jl b/core/src/model.jl
index 9ad1f0730..5a5e629fd 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -104,12 +104,10 @@ function Model(config::Config)::Model
     @assert eps(t_end) < 3600 "Simulation time too long"
     t0 = zero(t_end)
     timespan = (t0, t_end)
-    saveat = convert_saveat(config.solver.saveat, t_end)
+    saveat_state = convert_saveat(config.solver.saveat, t_end)
+    saveat_flow = convert_saveat(config.solver.saveat, t_end; state = false)
 
-    tstops_saveat =
-        isfinite(config.solver.saveat) ? range(t0, t_end; step = config.solver.saveat) :
-        Float64[]
-    tstops = sort(unique(vcat(tstops_flow_boundary, tstops_user, tstops_saveat)))
+    tstops = sort(unique(vcat(tstops_flow_boundary, tstops_user, saveat_flow)))
 
     jac_prototype = config.solver.sparse ? get_jac_prototype(parameters) : nothing
     RHS = ODEFunction(water_balance!; jac_prototype)
@@ -119,7 +117,7 @@ function Model(config::Config)::Model
     end
     @debug "Setup ODEProblem."
 
-    callback, saved = create_callbacks(parameters, config; saveat)
+    callback, saved = create_callbacks(parameters, config; saveat_flow)
     @debug "Created callbacks."
 
     # Initialize the integrator, providing all solver options as described in
@@ -135,7 +133,7 @@ function Model(config::Config)::Model
         callback,
         tstops,
         isoutofdomain = (u, p, t) -> any(<(0), u.storage),
-        saveat,
+        saveat = saveat_state,
         config.solver.adaptive,
         dt = something(config.solver.dt, t0),
         config.solver.dtmin,
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 110272fcb..6b6a6cb7d 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -479,13 +479,22 @@ end
 
 @testitem "mean_flow" begin
     toml_path =
-        normpath(@__DIR__, "../../generated_testmodels/linear_resistance/ribasim.toml")
+        normpath(@__DIR__, "../../generated_testmodels/flow_boundary_time/ribasim.toml")
     @test ispath(toml_path)
-    t_end = 3.16224e7
-
-    config = Ribasim.Config(toml_path; solver_saveat = t_end)
-    model = Ribasim.run(config)
-    mean_flow = model.saved.flow.saveval[1]
-    storage = Ribasim.get_storages_and_levels(model).storage[1, :]
-    @test storage[2] ≈ storage[1] - t_end * mean_flow[3] atol = 1
+    function get_flow(solver_adaptive::Bool, solver_saveat::Float64)::Vector{Float64}
+        config = Ribasim.Config(toml_path; solver_adaptive, solver_saveat)
+        model = Ribasim.run(config)
+        @show model.integrator.sol.t
+        df = DataFrame(Ribasim.flow_table(model))
+        return filter(
+            [:from_node_id, :to_node_id] => (from, to) -> from === 3 && to === 2,
+            df,
+        ).flow_rate
+    end
+    @test all(get_flow(true, 86400.0) .≈ 1.0) # broken
+    @test all(get_flow(true, 0.0) .≈ 1.0)
+    @test all(get_flow(true, Inf) .≈ 1.0)
+    @test all(get_flow(false, 86400.0) .≈ 1.0) # broken
+    @test all(get_flow(false, 0.0) .≈ 1.0)
+    @test all(get_flow(false, Inf) .≈ 1.0) # broken
 end

From ec01f499bef761692667b60038c00902a5d64d20 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 22 Feb 2024 18:07:51 +0100
Subject: [PATCH 07/19] Pass tests

---
 core/src/callback.jl         | 8 +++-----
 core/src/config.jl           | 2 +-
 core/src/graph.jl            | 4 ++--
 core/test/run_models_test.jl | 9 +++++----
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 18edbcf97..82b839dac 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -105,16 +105,16 @@ function integrate_flows!(u, t, integrator)::Nothing
     flow = get_tmp(flow, 0)
     flow_vertical = get_tmp(flow_vertical, 0)
     Δt = t - tprev
-    @show Δt
 
-    flow_effective = if length(integrator.sol.t) == 2
+    flow_effective = if isnan(flow_prev[1])
         # If flow_prev is not populated yet
         flow
     else
         0.5 * (flow + flow_prev)
     end
 
-    flow_vertical_effective = if length(integrator.sol.t) == 2
+    flow_vertical_effective = if isnan(flow_vertical_prev[1])
+        # If flow_vertical_prev is not populated yet
         flow_vertical
     else
         0.5 * (flow_vertical + flow_vertical_prev)
@@ -434,8 +434,6 @@ function save_flow(u, t, integrator)
         t_end - t >= saveat ? saveat : t - tprev
     end
 
-    @show Δt
-
     mean_flow_vertical = flow_vertical_integrated / Δt
     mean_flow = flow_integrated / Δt
 
diff --git a/core/src/config.jl b/core/src/config.jl
index f75845230..cf33aa91a 100644
--- a/core/src/config.jl
+++ b/core/src/config.jl
@@ -252,7 +252,7 @@ function convert_saveat(
         if state
             saveat
         else
-            collect(range(0.0, t_end; step = saveat))
+            collect(range(0.0, t_end; step = saveat))[2:end]
         end
     else
         @error "Invalid saveat" saveat
diff --git a/core/src/graph.jl b/core/src/graph.jl
index 9cfda9a80..feca6c51a 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -90,10 +90,10 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
     end
 
     flow = zeros(flow_counter)
-    flow_prev = zeros(flow_counter)
+    flow_prev = fill(NaN, flow_counter)
     flow_integrated = zeros(flow_counter)
     flow_vertical = zeros(flow_vertical_counter)
-    flow_vertical_prev = zeros(flow_vertical_counter)
+    flow_vertical_prev = fill(NaN, flow_vertical_counter)
     flow_vertical_integrated = zeros(flow_vertical_counter)
     if config.solver.autodiff
         flow = DiffCache(flow, chunk_sizes)
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 6b6a6cb7d..122e452d6 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -478,23 +478,24 @@ end
 end
 
 @testitem "mean_flow" begin
+    using DataFrames: DataFrame
+
     toml_path =
         normpath(@__DIR__, "../../generated_testmodels/flow_boundary_time/ribasim.toml")
     @test ispath(toml_path)
     function get_flow(solver_adaptive::Bool, solver_saveat::Float64)::Vector{Float64}
         config = Ribasim.Config(toml_path; solver_adaptive, solver_saveat)
         model = Ribasim.run(config)
-        @show model.integrator.sol.t
         df = DataFrame(Ribasim.flow_table(model))
         return filter(
             [:from_node_id, :to_node_id] => (from, to) -> from === 3 && to === 2,
             df,
         ).flow_rate
     end
-    @test all(get_flow(true, 86400.0) .≈ 1.0) # broken
+    @test all(get_flow(true, 86400.0) .≈ 1.0)
     @test all(get_flow(true, 0.0) .≈ 1.0)
     @test all(get_flow(true, Inf) .≈ 1.0)
-    @test all(get_flow(false, 86400.0) .≈ 1.0) # broken
+    @test all(get_flow(false, 86400.0) .≈ 1.0)
     @test all(get_flow(false, 0.0) .≈ 1.0)
-    @test all(get_flow(false, Inf) .≈ 1.0) # broken
+    @test all(get_flow(false, Inf) .≈ 1.0)
 end

From 51175e6a57418e473f06af821b06e366d5cbe40b Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 22 Feb 2024 20:39:53 +0100
Subject: [PATCH 08/19] Pass old tests as well

---
 core/src/callback.jl                          | 21 +++++++++------
 core/src/config.jl                            |  2 +-
 core/src/model.jl                             |  2 +-
 core/test/run_models_test.jl                  |  2 +-
 .../ribasim_testmodels/discrete_control.py    | 26 +++++++++----------
 5 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 82b839dac..b6d4e00cd 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -32,6 +32,7 @@ function create_callbacks(
     parameters::Parameters,
     config::Config;
     saveat_flow,
+    saveat_state,
 )::Tuple{CallbackSet, SavedResults}
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
     callbacks = SciMLBase.DECallback[]
@@ -68,7 +69,7 @@ function create_callbacks(
         export_cb = SavingCallback(
             save_subgrid_level,
             saved_subgrid_level;
-            saveat,
+            saveat = saveat_state,
             save_start = true,
         )
         push!(callbacks, export_cb)
@@ -91,6 +92,9 @@ function create_callbacks(
     return callback, saved
 end
 
+"""
+Integrate flows over timesteps
+"""
 function integrate_flows!(u, t, integrator)::Nothing
     (; p, tprev) = integrator
     (; graph) = p
@@ -106,19 +110,20 @@ function integrate_flows!(u, t, integrator)::Nothing
     flow_vertical = get_tmp(flow_vertical, 0)
     Δt = t - tprev
 
-    flow_effective = if isnan(flow_prev[1])
+    flow_effective = if !isempty(flow_prev) && isnan(flow_prev[1])
         # If flow_prev is not populated yet
         flow
     else
         0.5 * (flow + flow_prev)
     end
 
-    flow_vertical_effective = if isnan(flow_vertical_prev[1])
-        # If flow_vertical_prev is not populated yet
-        flow_vertical
-    else
-        0.5 * (flow_vertical + flow_vertical_prev)
-    end
+    flow_vertical_effective =
+        if !isempty(flow_vertical_prev) && isnan(flow_vertical_prev[1])
+            # If flow_vertical_prev is not populated yet
+            flow_vertical
+        else
+            0.5 * (flow_vertical + flow_vertical_prev)
+        end
 
     @. flow_integrated += flow_effective * Δt
     @. flow_vertical_integrated += flow_vertical_effective * Δt
diff --git a/core/src/config.jl b/core/src/config.jl
index cf33aa91a..1ed831a11 100644
--- a/core/src/config.jl
+++ b/core/src/config.jl
@@ -252,7 +252,7 @@ function convert_saveat(
         if state
             saveat
         else
-            collect(range(0.0, t_end; step = saveat))[2:end]
+            collect(range(saveat, t_end; step = saveat))
         end
     else
         @error "Invalid saveat" saveat
diff --git a/core/src/model.jl b/core/src/model.jl
index 5a5e629fd..c7396833a 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -117,7 +117,7 @@ function Model(config::Config)::Model
     end
     @debug "Setup ODEProblem."
 
-    callback, saved = create_callbacks(parameters, config; saveat_flow)
+    callback, saved = create_callbacks(parameters, config; saveat_flow, saveat_state)
     @debug "Created callbacks."
 
     # Initialize the integrator, providing all solver options as described in
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 122e452d6..52187a4d7 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -193,7 +193,7 @@ end
     @test model.integrator.sol.u[end] ≈ Float32[519.8817, 519.8798, 339.3959, 1418.4331] skip =
         Sys.isapple() atol = 1.5
 
-    @test length(logger.logs) == 8
+    @test length(logger.logs) == 11
     @test logger.logs[1].level == Debug
     @test logger.logs[1].message == "Read database into memory."
 
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 324e34949..e6554ec71 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -69,7 +69,18 @@ def pump_discrete_control_model() -> ribasim.Model:
 
     state = pd.DataFrame(data={"node_id": [1, 3], "level": [1.0, 1e-5]})
 
-    basin = ribasim.Basin(profile=profile, state=state)
+    static = pd.DataFrame(
+        data={
+            "node_id": [3],
+            "drainage": [0.0],
+            "potential_evaporation": [0.0],
+            "infiltration": [0.0],
+            "precipitation": [1e-9],
+            "urban_runoff": [0.0],
+        }
+    )
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Setup the discrete control:
     condition = pd.DataFrame(
@@ -325,20 +336,9 @@ def level_boundary_condition_model():
         }
     )
 
-    static = pd.DataFrame(
-        data={
-            "node_id": [3],
-            "drainage": [0.0],
-            "potential_evaporation": [0.0],
-            "infiltration": [0.0],
-            "precipitation": [0.0],
-            "urban_runoff": [0.0],
-        }
-    )
-
     state = pd.DataFrame(data={"node_id": [3], "level": [2.5]})
 
-    basin = ribasim.Basin(profile=profile, static=static, state=state)
+    basin = ribasim.Basin(profile=profile, state=state)
 
     # Setup level boundary:
     level_boundary = ribasim.LevelBoundary(

From 59627572a1f07232630644c97b1cda2783c855ac Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 22 Feb 2024 20:56:31 +0100
Subject: [PATCH 09/19] timesteps -> tstops

---
 core/src/model.jl            |  6 +++---
 core/src/write.jl            |  2 +-
 core/test/allocation_test.jl |  2 +-
 core/test/control_test.jl    | 36 +++++++++++++++++-------------------
 core/test/equations_test.jl  | 10 +++++-----
 core/test/run_models_test.jl | 14 +++++++-------
 docs/core/modelconcept.qmd   |  2 +-
 docs/core/numerics.qmd       |  2 +-
 docs/core/usage.qmd          |  4 ++--
 9 files changed, 38 insertions(+), 40 deletions(-)

diff --git a/core/src/model.jl b/core/src/model.jl
index c7396833a..f72a97e83 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -155,17 +155,17 @@ function Model(config::Config)::Model
 end
 
 "Get all saved times in seconds since start"
-timesteps(model::Model)::Vector{Float64} = model.integrator.sol.t
+tstops(model::Model)::Vector{Float64} = model.integrator.sol.t
 
 "Get all saved times as a Vector{DateTime}"
 function datetimes(model::Model)::Vector{DateTime}
-    return datetime_since.(timesteps(model), model.config.starttime)
+    return datetime_since.(tstops(model), model.config.starttime)
 end
 
 function Base.show(io::IO, model::Model)
     (; config, integrator) = model
     t = datetime_since(integrator.t, config.starttime)
-    nsaved = length(timesteps(model))
+    nsaved = length(tstops(model))
     println(io, "Model(ts: $nsaved, t: $t)")
 end
 
diff --git a/core/src/write.jl b/core/src/write.jl
index 17f21a8bf..0dcafa59a 100644
--- a/core/src/write.jl
+++ b/core/src/write.jl
@@ -64,7 +64,7 @@ function get_storages_and_levels(
     (; sol, p) = integrator
 
     node_id = p.basin.node_id.values::Vector{NodeID}
-    tsteps = datetime_since.(timesteps(model), config.starttime)
+    tsteps = datetime_since.(tstops(model), config.starttime)
 
     storage = hcat([collect(u_.storage) for u_ in sol.u]...)
     level = zero(storage)
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 9aa5e74f6..4a0bf734f 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -330,7 +330,7 @@ end
     model = Ribasim.run(toml_path)
 
     storage = Ribasim.get_storages_and_levels(model).storage[1, :]
-    t = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
 
     p = model.integrator.p
     (; user, graph, allocation, basin, target_level) = p
diff --git a/core/test/control_test.jl b/core/test/control_test.jl
index 2a76aa402..af4bd4f11 100644
--- a/core/test/control_test.jl
+++ b/core/test/control_test.jl
@@ -32,15 +32,15 @@
           ["off", "active", "on", "off", "inactive"]
 
     level = Ribasim.get_storages_and_levels(model).level
-    timesteps = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
 
     # Control times
     t_1 = discrete_control.record.time[3]
-    t_1_index = findfirst(timesteps .≈ t_1)
+    t_1_index = findfirst(t .≈ t_1)
     @test level[1, t_1_index] ≈ discrete_control.greater_than[1]
 
     t_2 = discrete_control.record.time[4]
-    t_2_index = findfirst(timesteps .≈ t_2)
+    t_2_index = findfirst(t .≈ t_2)
     @test level[2, t_2_index] ≈ discrete_control.greater_than[2]
 
     flow = get_tmp(graph[].flow, 0)
@@ -56,9 +56,9 @@ end
 
     Δt = discrete_control.look_ahead[1]
 
-    timesteps = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     t_control = discrete_control.record.time[2]
-    t_control_index = searchsortedfirst(timesteps, t_control)
+    t_control_index = searchsortedfirst(t, t_control)
 
     greater_than = discrete_control.greater_than[1]
     flow_t_control = flow_boundary.flow_rate[1](t_control)
@@ -80,9 +80,9 @@ end
 
     Δt = discrete_control.look_ahead[1]
 
-    timesteps = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     t_control = discrete_control.record.time[2]
-    t_control_index = searchsortedfirst(timesteps, t_control)
+    t_control_index = searchsortedfirst(t, t_control)
 
     greater_than = discrete_control.greater_than[1]
     level_t_control = level_boundary.level[1](t_control)
@@ -100,11 +100,11 @@ end
     (; basin, pid_control, flow_boundary) = p
 
     level = Ribasim.get_storages_and_levels(model).level[1, :]
-    timesteps = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
 
     target_itp = pid_control.target[1]
     t_target_change = target_itp.t[2]
-    idx_target_change = searchsortedlast(timesteps, t_target_change)
+    idx_target_change = searchsortedlast(t, t_target_change)
 
     K_p, K_i, _ = pid_control.pid_params[2](0)
     target_level = pid_control.target[2](0)
@@ -119,15 +119,13 @@ end
     phi = atan(du0 / (A * Δlevel) - alpha) / omega
     a = abs(Δlevel / cos(phi))
     # This bound is the exact envelope of the analytical solution
-    bound = @. a * exp(alpha * timesteps[1:idx_target_change])
+    bound = @. a * exp(alpha * t[1:idx_target_change])
     eps = 5e-3
     # Initial convergence to target level
     @test all(@. abs(level[1:idx_target_change] - target_level) < bound + eps)
     # Later closeness to target level
     @test all(
-        @. abs(
-            level[idx_target_change:end] - target_itp(timesteps[idx_target_change:end]),
-        ) < 5e-2
+        @. abs(level[idx_target_change:end] - target_itp(t[idx_target_change:end])) < 5e-2
     )
 end
 
@@ -163,7 +161,7 @@ end
     (; discrete_control) = p
     (; record, greater_than) = discrete_control
     level = Ribasim.get_storages_and_levels(model).level[1, :]
-    timesteps = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
 
     t_none_1 = discrete_control.record.time[2]
     t_in = discrete_control.record.time[3]
@@ -172,9 +170,9 @@ end
     level_min = greater_than[1]
     setpoint = greater_than[2]
 
-    t_1_none_index = findfirst(timesteps .≈ t_none_1)
-    t_in_index = findfirst(timesteps .≈ t_in)
-    t_2_none_index = findfirst(timesteps .≈ t_none_2)
+    t_1_none_index = findfirst(t .≈ t_none_1)
+    t_in_index = findfirst(t .≈ t_in)
+    t_2_none_index = findfirst(t .≈ t_none_2)
 
     @test record.control_state == ["out", "none", "in", "none"]
     @test level[t_1_none_index] ≈ setpoint
@@ -194,7 +192,7 @@ end
     p = model.integrator.p
     (; discrete_control, pid_control) = p
 
-    timesteps = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     level = Ribasim.get_storages_and_levels(model).level[1, :]
 
     target_high =
@@ -203,7 +201,7 @@ end
         pid_control.control_mapping[(NodeID(:PidControl, 6), "target_low")].target.u[1]
 
     t_target_jump = discrete_control.record.time[2]
-    t_idx_target_jump = searchsortedlast(timesteps, t_target_jump)
+    t_idx_target_jump = searchsortedlast(t, t_target_jump)
 
     @test isapprox(level[t_idx_target_jump], target_high, atol = 1e-1)
     @test isapprox(level[end], target_low, atol = 1e-1)
diff --git a/core/test/equations_test.jl b/core/test/equations_test.jl
index 7d00fb454..07099467e 100644
--- a/core/test/equations_test.jl
+++ b/core/test/equations_test.jl
@@ -25,7 +25,7 @@
     @test successful_retcode(model)
     (; p) = model.integrator
 
-    t = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     storage = Ribasim.get_storages_and_levels(model).storage[1, :]
     A = p.basin.area[1][2]  # needs to be constant
     u0 = A * 10.0
@@ -57,7 +57,7 @@ end
     @test successful_retcode(model)
     p = model.integrator.p
 
-    t = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     storage = Ribasim.get_storages_and_levels(model).storage[1, :]
     basin_area = p.basin.area[1][2]
     storage_min = 50.005
@@ -93,7 +93,7 @@ end
     p = model.integrator.p
     (; manning_resistance) = p
 
-    t = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     storage_both = Ribasim.get_storages_and_levels(model).storage
     storage = storage_both[1, :]
     storage_min = 50.005
@@ -131,7 +131,7 @@ end
     (; basin, pid_control) = p
 
     storage = Ribasim.get_storages_and_levels(model).storage[:]
-    t = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     SP = pid_control.target[1](0)
     K_p, K_i, K_d = pid_control.pid_params[1](0)
 
@@ -182,7 +182,7 @@ end
     frac = fractional_flow.fraction[1]
 
     storage_both = Ribasim.get_storages_and_levels(model).storage
-    t = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
 
     @test storage_both[1, :] ≈ @. storage_both[1, 1] + t * (frac * q_boundary - q_pump)
     @test storage_both[2, :] ≈ @. storage_both[2, 1] + t * q_pump
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 52187a4d7..e0e5e0004 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -4,7 +4,7 @@
     using Tables.DataAPI: nrow
     using Dates: DateTime
     import Arrow
-    using Ribasim: timesteps
+    using Ribasim: get_tstops, tstops
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/trivial/ribasim.toml")
     @test ispath(toml_path)
@@ -84,7 +84,7 @@
     end
 
     @testset "Results size" begin
-        nsaved = length(timesteps(model))
+        nsaved = length(tstops(model))
         @test nsaved > 10
         # t0 has no flow, 2 flow edges and 2 boundary condition flows
         @test nrow(flow) == (nsaved - 1) * 4
@@ -189,7 +189,7 @@ end
     @test all(isconcretetype, fieldtypes(typeof(p)))
 
     @test successful_retcode(model)
-    @test allunique(Ribasim.timesteps(model))
+    @test allunique(Ribasim.tstops(model))
     @test model.integrator.sol.u[end] ≈ Float32[519.8817, 519.8798, 339.3959, 1418.4331] skip =
         Sys.isapple() atol = 1.5
 
@@ -355,7 +355,7 @@ end
     (; level) = level_boundary
     level = level[1]
 
-    timesteps = model.saved.flow.t
+    t = model.saved.flow.t
     flow = DataFrame(Ribasim.flow_table(model))
     outlet_flow =
         filter([:from_node_id, :to_node_id] => (from, to) -> from === 2 && to === 3, flow)
@@ -364,14 +364,14 @@ end
         level.t[2] * (outlet.min_crest_level[1] - level.u[1]) / (level.u[2] - level.u[1])
 
     # No outlet flow when upstream level is below minimum crest level
-    @test all(@. outlet_flow.flow_rate[timesteps <= t_min_crest_level] == 0)
+    @test all(@. outlet_flow.flow_rate[t <= t_min_crest_level] == 0)
 
-    timesteps = Ribasim.timesteps(model)
+    t = Ribasim.tstops(model)
     t_maximum_level = level.t[2]
     level_basin = Ribasim.get_storages_and_levels(model).level[:]
 
     # Basin level converges to stable level boundary level
-    all(isapprox.(level_basin[timesteps .>= t_maximum_level], level.u[3], atol = 5e-2))
+    all(isapprox.(level_basin[t .>= t_maximum_level], level.u[3], atol = 5e-2))
 end
 
 @testitem "User" begin
diff --git a/docs/core/modelconcept.qmd b/docs/core/modelconcept.qmd
index 8f13075db..f781795d2 100644
--- a/docs/core/modelconcept.qmd
+++ b/docs/core/modelconcept.qmd
@@ -41,7 +41,7 @@ Ribasim is essentially a continuous model, rather than daily or hourly.
 If you want to use hourly forcing, you only need to make sure that your forcing data contains hourly updates.
 The output frequency can be configured independently.
 To be able to write a closed water balance, we accumulate the fluxes.
-This way any variations in between timesteps are also included, and we can output in `m³` rather than `m³s⁻¹`.
+This way any variations in between tstops are also included, and we can output in `m³` rather than `m³s⁻¹`.
 
 ## Space {#sec-space}
 
diff --git a/docs/core/numerics.qmd b/docs/core/numerics.qmd
index e54f69928..b20228d06 100644
--- a/docs/core/numerics.qmd
+++ b/docs/core/numerics.qmd
@@ -106,4 +106,4 @@ Rather, the basin gets a very small storage.
 The reduction factors help with performance, but are also an important tool to avoid getting negative storage in basins.
 Negative storage needs to be avoided since it is not a real solution, and would introduce water into the model that doesn't exist.
 Another tool used to avoid negative storage is the [`isoutoutofdomain`](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/) option, which Ribasim makes use of.
-This reject timesteps that lead to negative storage, instead retrying with a smaller timestep.
+This reject tstops that lead to negative storage, instead retrying with a smaller timestep.
diff --git a/docs/core/usage.qmd b/docs/core/usage.qmd
index aed51e8be..65ad6ae97 100644
--- a/docs/core/usage.qmd
+++ b/docs/core/usage.qmd
@@ -273,8 +273,8 @@ to have a reasonable and safe default, a value must be provided in the static ta
 This table is the transient form of the `Basin` table.
 The only difference is that a time column is added.
 The table must by sorted by time, and per time it must be sorted by `node_id`.
-A linear interpolation between the given timesteps is currently done if the
-solver takes timesteps between the given data points. More options will be available later.
+A linear interpolation between the given tstops is currently done if the
+solver takes tstops between the given data points. More options will be available later.
 
 ## Basin / state
 

From f757c37addd0a4361da129ae1cb9976a2e270a1d Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 26 Feb 2024 14:11:31 +0100
Subject: [PATCH 10/19] Use dt from integrator

---
 core/src/callback.jl        | 20 ++++++++++++--------
 core/src/graph.jl           |  6 ++++--
 core/src/solve.jl           |  2 --
 core/test/equations_test.jl |  6 +++---
 4 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 853b97a3b..761d89526 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -96,7 +96,7 @@ end
 Integrate flows over timesteps
 """
 function integrate_flows!(u, t, integrator)::Nothing
-    (; p, tprev) = integrator
+    (; p, dt) = integrator
     (; graph) = p
     (;
         flow,
@@ -108,7 +108,6 @@ function integrate_flows!(u, t, integrator)::Nothing
     ) = graph[]
     flow = get_tmp(flow, 0)
     flow_vertical = get_tmp(flow_vertical, 0)
-    Δt = t - tprev
 
     flow_effective = if !isempty(flow_prev) && isnan(flow_prev[1])
         # If flow_prev is not populated yet
@@ -125,8 +124,8 @@ function integrate_flows!(u, t, integrator)::Nothing
             0.5 * (flow_vertical + flow_vertical_prev)
         end
 
-    @. flow_integrated += flow_effective * Δt
-    @. flow_vertical_integrated += flow_vertical_effective * Δt
+    @. flow_integrated += flow_effective * dt
+    @. flow_vertical_integrated += flow_vertical_effective * dt
 
     copyto!(flow_prev, flow)
     copyto!(flow_vertical_prev, flow_vertical)
@@ -425,18 +424,23 @@ end
 
 "Copy the current flow to the SavedValues"
 function save_flow(u, t, integrator)
-    (; tprev, p) = integrator
+    (; dt, p) = integrator
     (; graph) = p
     (; flow_integrated, flow_vertical_integrated, saveat) = graph[]
 
     Δt = if iszero(saveat)
-        t - tprev
+        dt
     elseif isinf(saveat)
         t
     else
         t_end = integrator.sol.prob.tspan[2]
-        # The last interval might be shorter than saveat
-        t_end - t >= saveat ? saveat : t - tprev
+        if t_end - t > saveat
+            saveat
+        else
+            # The last interval might be shorter than saveat
+            rem = t % saveat
+            iszero(rem) ? saveat : rem
+        end
     end
 
     mean_flow_vertical = flow_vertical_integrated / Δt
diff --git a/core/src/graph.jl b/core/src/graph.jl
index d39a70fdd..6ebf2a8a4 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -6,7 +6,8 @@ and data of edges (EdgeMetadata):
 [`EdgeMetadata`](@ref)
 """
 function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGraph
-    node_rows = execute(db, "SELECT node_id, node_type, subnetwork_id FROM Node ORDER BY fid")
+    node_rows =
+        execute(db, "SELECT node_id, node_type, subnetwork_id FROM Node ORDER BY fid")
     edge_rows = execute(
         db,
         "SELECT fid, from_node_type, from_node_id, to_node_type, to_node_id, edge_type, subnetwork_id FROM Edge ORDER BY fid",
@@ -44,7 +45,8 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
             end
             push!(node_ids[allocation_network_id], node_id)
         end
-        graph[node_id] = NodeMetadata(Symbol(snake_case(row.node_type)), allocation_network_id)
+        graph[node_id] =
+            NodeMetadata(Symbol(snake_case(row.node_type)), allocation_network_id)
         if row.node_type in nonconservative_nodetypes
             flow_vertical_counter += 1
             flow_vertical_dict[node_id] = flow_vertical_counter
diff --git a/core/src/solve.jl b/core/src/solve.jl
index d00b0115e..b090dfa26 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -622,8 +622,6 @@ function formulate_du!(
     basin::Basin,
     storage::AbstractVector,
 )::Nothing
-    (; flow_vertical) = graph[]
-    flow_vertical = get_tmp(flow_vertical, storage)
     # loop over basins
     # subtract all outgoing flows
     # add all ingoing flows
diff --git a/core/test/equations_test.jl b/core/test/equations_test.jl
index 59b52156a..ff3f61aa6 100644
--- a/core/test/equations_test.jl
+++ b/core/test/equations_test.jl
@@ -173,7 +173,9 @@ end
     toml_path = normpath(@__DIR__, "../../generated_testmodels/misc_nodes/ribasim.toml")
     @test ispath(toml_path)
     config = Ribasim.Config(toml_path)
-    model = Ribasim.run(config)
+    model = Ribasim.BMI.initialize(Ribasim.Model, toml_path)
+    @test config.solver.dt === model.integrator.dt
+    Ribasim.solve!(model)
     @test successful_retcode(model)
     p = model.integrator.p
     (; flow_boundary, fractional_flow, pump) = p
@@ -186,8 +188,6 @@ end
     storage_both = get_storages_and_levels(model).storage
     t = tstops(model)
     tspan = model.integrator.sol.prob.tspan
-
-    @test config.solver.dt === model.integrator.dt
     @test t ≈ range(tspan...; step = config.solver.saveat)
     @test storage_both[1, :] ≈ @. storage_both[1, 1] + t * (frac * q_boundary - q_pump)
     @test storage_both[2, :] ≈ @. storage_both[2, 1] + t * q_pump

From 2f027fa49328220eb51f189045bcfb69db758e04 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 26 Feb 2024 14:23:45 +0100
Subject: [PATCH 11/19] tstops -> tsaves

---
 core/src/callback.jl         |  5 +++--
 core/src/model.jl            |  2 +-
 core/src/write.jl            |  2 +-
 core/test/allocation_test.jl |  2 +-
 core/test/control_test.jl    | 12 ++++++------
 core/test/run_models_test.jl |  8 ++++----
 6 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 761d89526..215bdace9 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -93,7 +93,7 @@ function create_callbacks(
 end
 
 """
-Integrate flows over timesteps
+Integrate flows over the last timestep
 """
 function integrate_flows!(u, t, integrator)::Nothing
     (; p, dt) = integrator
@@ -422,7 +422,8 @@ function set_control_params!(p::Parameters, node_id::NodeID, control_state::Stri
     end
 end
 
-"Copy the current flow to the SavedValues"
+"Compute the average flows over the last saveat interval and write
+them to SavedValues"
 function save_flow(u, t, integrator)
     (; dt, p) = integrator
     (; graph) = p
diff --git a/core/src/model.jl b/core/src/model.jl
index ffd8b00fe..81c8f337f 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -155,7 +155,7 @@ function Model(config::Config)::Model
 end
 
 "Get all saved times in seconds since start"
-tstops(model::Model)::Vector{Float64} = model.integrator.sol.t
+tsaves(model::Model)::Vector{Float64} = model.integrator.sol.t
 
 "Get all saved times as a Vector{DateTime}"
 function datetimes(model::Model)::Vector{DateTime}
diff --git a/core/src/write.jl b/core/src/write.jl
index 0dcafa59a..9c7b39610 100644
--- a/core/src/write.jl
+++ b/core/src/write.jl
@@ -64,7 +64,7 @@ function get_storages_and_levels(
     (; sol, p) = integrator
 
     node_id = p.basin.node_id.values::Vector{NodeID}
-    tsteps = datetime_since.(tstops(model), config.starttime)
+    tsteps = datetime_since.(tsaves(model), config.starttime)
 
     storage = hcat([collect(u_.storage) for u_ in sol.u]...)
     level = zero(storage)
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index c6b407464..dd361397e 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -341,7 +341,7 @@ end
     model = Ribasim.run(toml_path)
 
     storage = Ribasim.get_storages_and_levels(model).storage[1, :]
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
 
     p = model.integrator.p
     (; user_demand, graph, allocation, basin, level_demand) = p
diff --git a/core/test/control_test.jl b/core/test/control_test.jl
index a4d916707..ae572e45e 100644
--- a/core/test/control_test.jl
+++ b/core/test/control_test.jl
@@ -32,7 +32,7 @@
           ["off", "active", "on", "off", "inactive"]
 
     level = Ribasim.get_storages_and_levels(model).level
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
 
     # Control times
     t_1 = discrete_control.record.time[3]
@@ -56,7 +56,7 @@ end
 
     Δt = discrete_control.look_ahead[1]
 
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     t_control = discrete_control.record.time[2]
     t_control_index = searchsortedfirst(t, t_control)
 
@@ -80,7 +80,7 @@ end
 
     Δt = discrete_control.look_ahead[1]
 
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     t_control = discrete_control.record.time[2]
     t_control_index = searchsortedfirst(t, t_control)
 
@@ -100,7 +100,7 @@ end
     (; basin, pid_control, flow_boundary) = p
 
     level = Ribasim.get_storages_and_levels(model).level[1, :]
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
 
     target_itp = pid_control.target[1]
     t_target_change = target_itp.t[2]
@@ -161,7 +161,7 @@ end
     (; discrete_control) = p
     (; record, greater_than) = discrete_control
     level = Ribasim.get_storages_and_levels(model).level[1, :]
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
 
     t_none_1 = discrete_control.record.time[2]
     t_in = discrete_control.record.time[3]
@@ -192,7 +192,7 @@ end
     p = model.integrator.p
     (; discrete_control, pid_control) = p
 
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     level = Ribasim.get_storages_and_levels(model).level[1, :]
 
     target_high =
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 3d96f5cb1..4d3fd0ae6 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -4,7 +4,7 @@
     using Tables.DataAPI: nrow
     using Dates: DateTime
     import Arrow
-    using Ribasim: get_tstops, tstops
+    using Ribasim: get_tstops, tsaves
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/trivial/ribasim.toml")
     @test ispath(toml_path)
@@ -84,7 +84,7 @@
     end
 
     @testset "Results size" begin
-        nsaved = length(tstops(model))
+        nsaved = length(tsaves(model))
         @test nsaved > 10
         # t0 has no flow, 2 flow edges and 2 boundary condition flows
         @test nrow(flow) == (nsaved - 1) * 4
@@ -189,7 +189,7 @@ end
     @test all(isconcretetype, fieldtypes(typeof(p)))
 
     @test successful_retcode(model)
-    @test allunique(Ribasim.tstops(model))
+    @test allunique(Ribasim.tsaves(model))
     @test model.integrator.sol.u[end] ≈ Float32[519.8817, 519.8798, 339.3959, 1418.4331] skip =
         Sys.isapple() atol = 1.5
 
@@ -366,7 +366,7 @@ end
     # No outlet flow when upstream level is below minimum crest level
     @test all(@. outlet_flow.flow_rate[t <= t_min_crest_level] == 0)
 
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     t_maximum_level = level.t[2]
     level_basin = Ribasim.get_storages_and_levels(model).level[:]
 

From 5601282dfd83c62028fb74a22fb02201fbab9756 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 26 Feb 2024 14:30:51 +0100
Subject: [PATCH 12/19] Fix tests

---
 core/src/model.jl           |  4 ++--
 core/test/equations_test.jl | 12 ++++++------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/core/src/model.jl b/core/src/model.jl
index 81c8f337f..6ca49230d 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -159,13 +159,13 @@ tsaves(model::Model)::Vector{Float64} = model.integrator.sol.t
 
 "Get all saved times as a Vector{DateTime}"
 function datetimes(model::Model)::Vector{DateTime}
-    return datetime_since.(tstops(model), model.config.starttime)
+    return datetime_since.(tsaves(model), model.config.starttime)
 end
 
 function Base.show(io::IO, model::Model)
     (; config, integrator) = model
     t = datetime_since(integrator.t, config.starttime)
-    nsaved = length(tstops(model))
+    nsaved = length(tsaves(model))
     println(io, "Model(ts: $nsaved, t: $t)")
 end
 
diff --git a/core/test/equations_test.jl b/core/test/equations_test.jl
index ff3f61aa6..9723a178e 100644
--- a/core/test/equations_test.jl
+++ b/core/test/equations_test.jl
@@ -25,7 +25,7 @@
     @test successful_retcode(model)
     (; p) = model.integrator
 
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     storage = Ribasim.get_storages_and_levels(model).storage[1, :]
     A = p.basin.area[1][2]  # needs to be constant
     u0 = A * 10.0
@@ -57,7 +57,7 @@ end
     @test successful_retcode(model)
     p = model.integrator.p
 
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     storage = Ribasim.get_storages_and_levels(model).storage[1, :]
     basin_area = p.basin.area[1][2]
     storage_min = 50.005
@@ -93,7 +93,7 @@ end
     p = model.integrator.p
     (; manning_resistance) = p
 
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     storage_both = Ribasim.get_storages_and_levels(model).storage
     storage = storage_both[1, :]
     storage_min = 50.005
@@ -131,7 +131,7 @@ end
     (; basin, pid_control) = p
 
     storage = Ribasim.get_storages_and_levels(model).storage[:]
-    t = Ribasim.tstops(model)
+    t = Ribasim.tsaves(model)
     SP = pid_control.target[1](0)
     K_p, K_i, K_d = pid_control.pid_params[1](0)
 
@@ -168,7 +168,7 @@ end
 @testitem "MiscellaneousNodes" begin
     using PreallocationTools: get_tmp
     using SciMLBase: successful_retcode
-    using Ribasim: tstops, get_storages_and_levels
+    using Ribasim: tsaves, get_storages_and_levels
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/misc_nodes/ribasim.toml")
     @test ispath(toml_path)
@@ -186,7 +186,7 @@ end
     frac = fractional_flow.fraction[1]
 
     storage_both = get_storages_and_levels(model).storage
-    t = tstops(model)
+    t = tsaves(model)
     tspan = model.integrator.sol.prob.tspan
     @test t ≈ range(tspan...; step = config.solver.saveat)
     @test storage_both[1, :] ≈ @. storage_both[1, 1] + t * (frac * q_boundary - q_pump)

From 72e2b5816c062d4997d94ddab443a70f518a878f Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 26 Feb 2024 14:41:36 +0100
Subject: [PATCH 13/19] update docs

---
 docs/core/usage.qmd | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs/core/usage.qmd b/docs/core/usage.qmd
index 15791d287..9330874c5 100644
--- a/docs/core/usage.qmd
+++ b/docs/core/usage.qmd
@@ -680,8 +680,7 @@ The table is sorted by time, and per time it is sorted by `node_id`.
 
 ## Flow - `flow.arrow`
 
-The flow table contains results of the flow on every edge in the model, for each solver
-timestep.
+The flow table contains results of the flow on every edge in the model, averaged over each `saveat` interval.
 
 column         | type                | unit
 -------------- | ------------------- | ----

From b20c81465b63624a0e0b25cb4673d8feb484d3d9 Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Mon, 26 Feb 2024 22:46:10 +0100
Subject: [PATCH 14/19] Set timestamps at the beginning of the period

---
 core/src/write.jl            | 5 ++++-
 core/test/run_models_test.jl | 4 ++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/core/src/write.jl b/core/src/write.jl
index 9c7b39610..d88ac7a74 100644
--- a/core/src/write.jl
+++ b/core/src/write.jl
@@ -148,7 +148,10 @@ function flow_table(
     nflow = length(unique_edge_ids_flow)
     ntsteps = length(t)
 
-    time = repeat(datetime_since.(t, config.starttime); inner = nflow)
+    # the timestamp should represent the start of the period, not the end
+    t_starts = circshift(t, 1)
+    t_starts[1] = 0.0
+    time = repeat(datetime_since.(t_starts, config.starttime); inner = nflow)
     edge_id = repeat(unique_edge_ids_flow; outer = ntsteps)
     from_node_type = repeat(from_node_type; outer = ntsteps)
     from_node_id = repeat(from_node_id; outer = ntsteps)
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 4d3fd0ae6..2b2d4bb74 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -95,7 +95,7 @@
     end
 
     @testset "Results values" begin
-        @test flow.time[1] > DateTime(2020)
+        @test flow.time[1] == DateTime(2020)
         @test coalesce.(flow.edge_id[1:4], -1) == [-1, -1, 9, 11]
         @test flow.from_node_id[1:4] == [6, typemax(Int), 0, 6]
         @test flow.to_node_id[1:4] == [6, typemax(Int), typemax(Int), 0]
@@ -200,7 +200,7 @@ end
     table = Ribasim.flow_table(model)
 
     # flows are recorded at the end of each period, and are undefined at the start
-    @test unique(table.time) == Ribasim.datetimes(model)[2:end]
+    @test unique(table.time) == Ribasim.datetimes(model)[1:(end - 1)]
 
     # inflow = outflow over FractionalFlow
     t = table.time[1]

From a8b335c4bf61dbbe522da2ad8cf188267bf3ceda Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Tue, 27 Feb 2024 10:03:22 +0100
Subject: [PATCH 15/19] fix empty results

---
 core/src/write.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/core/src/write.jl b/core/src/write.jl
index d88ac7a74..b94ef0c36 100644
--- a/core/src/write.jl
+++ b/core/src/write.jl
@@ -150,7 +150,9 @@ function flow_table(
 
     # the timestamp should represent the start of the period, not the end
     t_starts = circshift(t, 1)
-    t_starts[1] = 0.0
+    if !isempty(t)
+        t_starts[1] = 0.0
+    end
     time = repeat(datetime_since.(t_starts, config.starttime); inner = nflow)
     edge_id = repeat(unique_edge_ids_flow; outer = ntsteps)
     from_node_type = repeat(from_node_type; outer = ntsteps)

From 2763867425a321c86250f6855b326b418c2b3fa7 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 27 Feb 2024 17:00:30 +0100
Subject: [PATCH 16/19] Adress most code comments

---
 core/src/callback.jl         | 47 ++++++++++++---------------
 core/src/config.jl           | 15 +++------
 core/src/model.jl            | 16 ++++++----
 core/test/run_models_test.jl | 62 ++++++++++++++++++++++++++++++------
 4 files changed, 86 insertions(+), 54 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 215bdace9..7b312f409 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -30,9 +30,8 @@ Returns the CallbackSet and the SavedValues for flow.
 """
 function create_callbacks(
     parameters::Parameters,
-    config::Config;
-    saveat_flow,
-    saveat_state,
+    config::Config,
+    saveat,
 )::Tuple{CallbackSet, SavedResults}
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
     callbacks = SciMLBase.DECallback[]
@@ -59,8 +58,14 @@ function create_callbacks(
 
     # save the flows over time, as a Vector of the nonzeros(flow)
     saved_flow = SavedValues(Float64, Vector{Float64})
-    save_flow_cb =
-        SavingCallback(save_flow, saved_flow; saveat = saveat_flow, save_start = false)
+    save_flow_cb = SavingCallback(
+        save_flow,
+        saved_flow;
+        # If saveat is a vector which contains 0.0 this callback will still be called
+        # at t = 0.0 despite save_start = false
+        saveat = saveat isa Vector ? filter(x -> x != 0.0, saveat) : saveat,
+        save_start = false,
+    )
     push!(callbacks, save_flow_cb)
 
     # interpolate the levels
@@ -69,7 +74,7 @@ function create_callbacks(
         export_cb = SavingCallback(
             save_subgrid_level,
             saved_subgrid_level;
-            saveat = saveat_state,
+            saveat,
             save_start = true,
         )
         push!(callbacks, export_cb)
@@ -109,23 +114,14 @@ function integrate_flows!(u, t, integrator)::Nothing
     flow = get_tmp(flow, 0)
     flow_vertical = get_tmp(flow_vertical, 0)
 
-    flow_effective = if !isempty(flow_prev) && isnan(flow_prev[1])
+    if !isempty(flow_prev) && isnan(flow_prev[1])
         # If flow_prev is not populated yet
-        flow
-    else
-        0.5 * (flow + flow_prev)
+        copyto!(flow_prev, flow)
+        copyto!(flow_vertical_prev, flow_vertical)
     end
 
-    flow_vertical_effective =
-        if !isempty(flow_vertical_prev) && isnan(flow_vertical_prev[1])
-            # If flow_vertical_prev is not populated yet
-            flow_vertical
-        else
-            0.5 * (flow_vertical + flow_vertical_prev)
-        end
-
-    @. flow_integrated += flow_effective * dt
-    @. flow_vertical_integrated += flow_vertical_effective * dt
+    @. flow_integrated += 0.5 * (flow + flow_prev) * dt
+    @. flow_vertical_integrated += 0.5 * (flow_vertical + flow_vertical_prev) * dt
 
     copyto!(flow_prev, flow)
     copyto!(flow_vertical_prev, flow_vertical)
@@ -317,8 +313,7 @@ function discrete_control_affect!(
 
     # What the local control state is
     # TODO: Check time elapsed since control change
-    control_state_now, control_state_start =
-        discrete_control.control_state[discrete_control_node_id]
+    control_state_now, _ = discrete_control.control_state[discrete_control_node_id]
 
     control_state_change = false
 
@@ -444,12 +439,12 @@ function save_flow(u, t, integrator)
         end
     end
 
-    mean_flow_vertical = flow_vertical_integrated / Δt
-    mean_flow = flow_integrated / Δt
-
+    mean_flow_all = vcat(flow_vertical_integrated, flow_integrated)
+    mean_flow_all ./= Δt
     fill!(flow_vertical_integrated, 0.0)
     fill!(flow_integrated, 0.0)
-    return vcat(mean_flow_vertical, mean_flow)
+
+    return mean_flow_all
 end
 
 function update_subgrid_level!(integrator)::Nothing
diff --git a/core/src/config.jl b/core/src/config.jl
index 7ecc8df68..07dc6afa4 100644
--- a/core/src/config.jl
+++ b/core/src/config.jl
@@ -232,27 +232,20 @@ function algorithm(solver::Solver)::OrdinaryDiffEqAlgorithm
 end
 
 "Convert the saveat Float64 from our Config to SciML's saveat"
-function convert_saveat(
-    saveat::Float64,
-    t_end::Float64,
-)::Tuple{Union{Float64, Vector{Float64}}, Vector{Float64}}
+function convert_saveat(saveat::Float64, t_end::Float64)::Union{Float64, Vector{Float64}}
     if iszero(saveat)
         # every step
-        saveat_state = Float64[]
-        saveat_flow = Float64[]
+        Float64[]
     elseif saveat == Inf
         # only the start and end
-        saveat_state = [0.0, t_end]
-        saveat_flow = [t_end]
+        [0.0, t_end]
     elseif isfinite(saveat)
         # every saveat seconds
-        saveat_state = saveat
-        saveat_flow = collect(range(saveat, t_end; step = saveat))
+        saveat
     else
         @error "Invalid saveat" saveat
         error("Invalid saveat")
     end
-    return saveat_state, saveat_flow
 end
 
 "Convert the dt from our Config to SciML stepsize control arguments"
diff --git a/core/src/model.jl b/core/src/model.jl
index 6ca49230d..009cd80c2 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -46,7 +46,7 @@ function Model(config::Config)::Model
     # All data from the database that we need during runtime is copied into memory,
     # so we can directly close it again.
     db = SQLite.DB(db_path)
-    local parameters, state, n, tstops, tstops_flow_boundary, tstops_user_demand
+    local parameters, state, n, tstops
     try
         parameters = Parameters(db, config)
 
@@ -76,9 +76,10 @@ function Model(config::Config)::Model
         # tell the solver to stop when new data comes in
         # TODO add all time tables here
         time_flow_boundary = load_structvector(db, config, FlowBoundaryTimeV1)
-        tstops_flow_boundary = get_tstops(time_flow_boundary.time, config.starttime)
+        tstops = Vector{Float64}[]
+        push!(tstops, get_tstops(time_flow_boundary.time, config.starttime))
         time_user_demand = load_structvector(db, config, UserDemandTimeV1)
-        tstops_user_demand = get_tstops(time_user_demand.time, config.starttime)
+        push!(tstops, get_tstops(time_user_demand.time, config.starttime))
 
         # use state
         state = load_structvector(db, config, BasinStateV1)
@@ -105,8 +106,9 @@ function Model(config::Config)::Model
     t0 = zero(t_end)
     timespan = (t0, t_end)
 
-    saveat_state, saveat_flow = convert_saveat(config.solver.saveat, t_end)
-    tstops = sort(unique(vcat(tstops_flow_boundary, tstops_user_demand, saveat_flow)))
+    saveat = convert_saveat(config.solver.saveat, t_end)
+    saveat isa Float64 && push!(tstops, range(0, t_end; step = saveat))
+    tstops = sort(unique(vcat(tstops...)))
     adaptive, dt = convert_dt(config.solver.dt)
 
     jac_prototype = config.solver.sparse ? get_jac_prototype(parameters) : nothing
@@ -117,7 +119,7 @@ function Model(config::Config)::Model
     end
     @debug "Setup ODEProblem."
 
-    callback, saved = create_callbacks(parameters, config; saveat_flow, saveat_state)
+    callback, saved = create_callbacks(parameters, config, saveat)
     @debug "Created callbacks."
 
     # Initialize the integrator, providing all solver options as described in
@@ -133,7 +135,7 @@ function Model(config::Config)::Model
         callback,
         tstops,
         isoutofdomain = (u, p, t) -> any(<(0), u.storage),
-        saveat = saveat_state,
+        saveat,
         adaptive,
         dt,
         config.solver.dtmin,
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 2b2d4bb74..c301b6116 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -483,21 +483,63 @@ end
     toml_path =
         normpath(@__DIR__, "../../generated_testmodels/flow_boundary_time/ribasim.toml")
     @test ispath(toml_path)
-    function get_flow(solver_adaptive::Bool, solver_saveat::Float64)::Vector{Float64}
-        config = Ribasim.Config(toml_path)
-        solver_dt = solver_adaptive ? nothing : config.solver.dt
+    function get_flow(solver_dt::Union{Float64, Nothing}, solver_saveat::Float64)
         config = Ribasim.Config(toml_path; solver_dt, solver_saveat)
         model = Ribasim.run(config)
         df = DataFrame(Ribasim.flow_table(model))
         return filter(
             [:from_node_id, :to_node_id] => (from, to) -> from === 3 && to === 2,
             df,
-        ).flow_rate
+        ).flow_rate,
+        Ribasim.tsaves(model)
     end
-    @test all(get_flow(true, 86400.0) .≈ 1.0)
-    @test all(get_flow(true, 0.0) .≈ 1.0)
-    @test all(get_flow(true, Inf) .≈ 1.0)
-    @test all(get_flow(false, 86400.0) .≈ 1.0)
-    @test all(get_flow(false, 0.0) .≈ 1.0)
-    @test all(get_flow(false, Inf) .≈ 1.0)
+
+    Δt = 24 * 24 * 60.0
+    t_end = 3.16224e7 # 366 days
+
+    # t_end % saveat = 0
+    saveat = 86400.0
+    flow, tstops = get_flow(nothing, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == t_end / saveat
+    @test length(tstops) == t_end / saveat + 1
+
+    flow, tstops = get_flow(Δt, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == t_end / saveat
+    @test length(tstops) == t_end / saveat + 1
+
+    # t_end % saveat != 0
+    saveat = round(10000 * π)
+    flow, tstops = get_flow(nothing, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == ceil(t_end / saveat)
+    @test length(tstops) == ceil(t_end / saveat) + 1
+
+    flow, tstops = get_flow(Δt, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == ceil(t_end / saveat)
+    @test length(tstops) == ceil(t_end / saveat) + 1
+
+    # Only save average over all flows in tspan
+    saveat = Inf
+    flow, tstops = get_flow(nothing, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == 1
+    @test length(tstops) == 2
+
+    flow, tstops = get_flow(Δt, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == 1
+    @test length(tstops) == 2
+
+    # Save all flows
+    saveat = 0.0
+    flow, tstops = get_flow(nothing, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == length(tstops) - 1
+
+    flow, tstops = get_flow(Δt, saveat)
+    @test all(flow .≈ 1.0)
+    @test length(flow) == length(tstops) - 1
 end

From 550920ea4dbd26f6af8417e635a4a023d001390f Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 28 Feb 2024 08:30:09 +0100
Subject: [PATCH 17/19] Final comments adressed

---
 core/src/config.jl          | 11 +++++++++--
 core/test/config_test.jl    |  1 +
 core/test/equations_test.jl |  2 +-
 docs/core/modelconcept.qmd  |  2 +-
 docs/core/numerics.qmd      |  2 +-
 docs/core/usage.qmd         |  6 +++---
 6 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/core/src/config.jl b/core/src/config.jl
index 07dc6afa4..c648cb83b 100644
--- a/core/src/config.jl
+++ b/core/src/config.jl
@@ -233,19 +233,26 @@ end
 
 "Convert the saveat Float64 from our Config to SciML's saveat"
 function convert_saveat(saveat::Float64, t_end::Float64)::Union{Float64, Vector{Float64}}
+    error = false
     if iszero(saveat)
         # every step
-        Float64[]
+        saveat = Float64[]
     elseif saveat == Inf
         # only the start and end
         [0.0, t_end]
     elseif isfinite(saveat)
         # every saveat seconds
         saveat
+        if saveat !== round(saveat)
+            @error "A finite saveat must be an integer number of seconds."
+        end
     else
+        error = true
         @error "Invalid saveat" saveat
-        error("Invalid saveat")
     end
+
+    error && error("Invalid saveat")
+    return saveat
 end
 
 "Convert the dt from our Config to SciML stepsize control arguments"
diff --git a/core/test/config_test.jl b/core/test/config_test.jl
index 535aebed7..786138738 100644
--- a/core/test/config_test.jl
+++ b/core/test/config_test.jl
@@ -71,6 +71,7 @@ end
     @test convert_saveat(Inf, t_end) == ([0.0, t_end], [t_end])
     @test_throws ErrorException convert_saveat(-Inf, t_end)
     @test_throws ErrorException convert_saveat(NaN, t_end)
+    @test_throws ErrorException convert_saveat(3.1415, t_end)
 
     @test convert_dt(nothing) == (true, 0.0)
     @test convert_dt(360.0) == (false, 360.0)
diff --git a/core/test/equations_test.jl b/core/test/equations_test.jl
index 9723a178e..95ca06d03 100644
--- a/core/test/equations_test.jl
+++ b/core/test/equations_test.jl
@@ -173,7 +173,7 @@ end
     toml_path = normpath(@__DIR__, "../../generated_testmodels/misc_nodes/ribasim.toml")
     @test ispath(toml_path)
     config = Ribasim.Config(toml_path)
-    model = Ribasim.BMI.initialize(Ribasim.Model, toml_path)
+    model = Ribasim.Model(toml_path)
     @test config.solver.dt === model.integrator.dt
     Ribasim.solve!(model)
     @test successful_retcode(model)
diff --git a/docs/core/modelconcept.qmd b/docs/core/modelconcept.qmd
index 386f41693..20f260399 100644
--- a/docs/core/modelconcept.qmd
+++ b/docs/core/modelconcept.qmd
@@ -41,7 +41,7 @@ Ribasim is essentially a continuous model, rather than daily or hourly.
 If you want to use hourly forcing, you only need to make sure that your forcing data contains hourly updates.
 The output frequency can be configured independently.
 To be able to write a closed water balance, we accumulate the fluxes.
-This way any variations in between tstops are also included, and we can output in `m³` rather than `m³s⁻¹`.
+This way any variations in between timesteps are also included, and we can output in `m³` rather than `m³s⁻¹`.
 
 ## Space {#sec-space}
 
diff --git a/docs/core/numerics.qmd b/docs/core/numerics.qmd
index b20228d06..ca25ca183 100644
--- a/docs/core/numerics.qmd
+++ b/docs/core/numerics.qmd
@@ -106,4 +106,4 @@ Rather, the basin gets a very small storage.
 The reduction factors help with performance, but are also an important tool to avoid getting negative storage in basins.
 Negative storage needs to be avoided since it is not a real solution, and would introduce water into the model that doesn't exist.
 Another tool used to avoid negative storage is the [`isoutoutofdomain`](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/) option, which Ribasim makes use of.
-This reject tstops that lead to negative storage, instead retrying with a smaller timestep.
+This rejects timesteps that lead to negative storage, instead retrying with a smaller timestep.
diff --git a/docs/core/usage.qmd b/docs/core/usage.qmd
index 9330874c5..85ce8f83a 100644
--- a/docs/core/usage.qmd
+++ b/docs/core/usage.qmd
@@ -271,8 +271,7 @@ to have a reasonable and safe default, a value must be provided in the static ta
 This table is the transient form of the `Basin` table.
 The only difference is that a time column is added.
 The table must by sorted by time, and per time it must be sorted by `node_id`.
-A linear interpolation between the given tstops is currently done if the
-solver takes tstops between the given data points. More options will be available later.
+At the given timestamps the values are set in the simulation, such that the timeseries can be seen as forward filled.
 
 ## Basin / state
 
@@ -680,7 +679,8 @@ The table is sorted by time, and per time it is sorted by `node_id`.
 
 ## Flow - `flow.arrow`
 
-The flow table contains results of the flow on every edge in the model, averaged over each `saveat` interval.
+The flow table contains calculated mean flows for every flow edge in the model.
+In the time column the start of the period is indicated.
 
 column         | type                | unit
 -------------- | ------------------- | ----

From 5f5e8de29dd69bccbd13661331e47ee4e0e5365f Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 28 Feb 2024 08:39:20 +0100
Subject: [PATCH 18/19] Fix tests

---
 core/src/config.jl       | 12 ++++++------
 core/test/config_test.jl |  7 +++----
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/core/src/config.jl b/core/src/config.jl
index c648cb83b..831c82106 100644
--- a/core/src/config.jl
+++ b/core/src/config.jl
@@ -233,25 +233,25 @@ end
 
 "Convert the saveat Float64 from our Config to SciML's saveat"
 function convert_saveat(saveat::Float64, t_end::Float64)::Union{Float64, Vector{Float64}}
-    error = false
+    errors = false
     if iszero(saveat)
         # every step
         saveat = Float64[]
     elseif saveat == Inf
         # only the start and end
-        [0.0, t_end]
+        saveat = [0.0, t_end]
     elseif isfinite(saveat)
         # every saveat seconds
-        saveat
         if saveat !== round(saveat)
-            @error "A finite saveat must be an integer number of seconds."
+            errors = true
+            @error "A finite saveat must be an integer number of seconds." saveat
         end
     else
-        error = true
+        errors = true
         @error "Invalid saveat" saveat
     end
 
-    error && error("Invalid saveat")
+    errors && error("Invalid saveat")
     return saveat
 end
 
diff --git a/core/test/config_test.jl b/core/test/config_test.jl
index 786138738..1c4aacf47 100644
--- a/core/test/config_test.jl
+++ b/core/test/config_test.jl
@@ -65,10 +65,9 @@ end
     algorithm(Solver(; algorithm = "Euler", autodiff = true))
 
     t_end = 100.0
-    @test convert_saveat(0.0, t_end) == (Float64[], Float64[])
-    @test convert_saveat(60.0, t_end) == (60.0, [60.0])
-    @test convert_saveat(Inf, t_end) == ([0.0, t_end], [t_end])
-    @test convert_saveat(Inf, t_end) == ([0.0, t_end], [t_end])
+    @test convert_saveat(0.0, t_end) == Float64[]
+    @test convert_saveat(60.0, t_end) == 60.0
+    @test convert_saveat(Inf, t_end) == [0.0, t_end]
     @test_throws ErrorException convert_saveat(-Inf, t_end)
     @test_throws ErrorException convert_saveat(NaN, t_end)
     @test_throws ErrorException convert_saveat(3.1415, t_end)

From f86410c1caa89ead171cc8d42fca87daa162a7b5 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 28 Feb 2024 11:19:36 +0100
Subject: [PATCH 19/19] Last comment adressed

---
 core/test/run_models_test.jl | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index c301b6116..8a1414275 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -487,11 +487,12 @@ end
         config = Ribasim.Config(toml_path; solver_dt, solver_saveat)
         model = Ribasim.run(config)
         df = DataFrame(Ribasim.flow_table(model))
-        return filter(
-            [:from_node_id, :to_node_id] => (from, to) -> from === 3 && to === 2,
-            df,
-        ).flow_rate,
-        Ribasim.tsaves(model)
+        flow =
+            filter(
+                [:from_node_id, :to_node_id] => (from, to) -> from === 3 && to === 2,
+                df,
+            ).flow_rate
+        flow, Ribasim.tsaves(model)
     end
 
     Δt = 24 * 24 * 60.0