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] 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