Mean output flows

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)
@@ -31,7 +31,8 @@ Returns the CallbackSet and the SavedValues for flow.
 function create_callbacks(
     parameters::Parameters,
     config::Config;
-    saveat,
+    saveat_flow,
+    saveat_state,
 )::Tuple{CallbackSet, SavedResults}
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
     callbacks = SciMLBase.DECallback[]
@@ -40,6 +41,9 @@ function create_callbacks(
     basin_cb = PresetTimeCallback(tstops, update_basin)
     push!(callbacks, basin_cb)
 
+    integrating_flows_cb = FunctionCallingCallback(integrate_flows!; func_start = false)
+    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)
@@ -55,7 +59,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
@@ -64,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)
@@ -87,6 +92,47 @@ function create_callbacks(
     return callback, saved
 end
 
+"""
+Integrate flows over timesteps
+"""
+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_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 !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
+
+    copyto!(flow_prev, flow)
+    copyto!(flow_vertical_prev, flow_vertical)
+    return nothing
+end
+
 """
 Listens for changes in condition truths.
 """
@@ -379,10 +425,26 @@ 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),
-    )
+    (; 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
+
+    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

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(saveat, t_end; step = saveat))
+        end
     else
         @error "Invalid saveat" saveat
         error("Invalid saveat")

core/src/graph.jl

@@ -90,7 +90,11 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
     end
 
     flow = zeros(flow_counter)
+    flow_prev = fill(NaN, flow_counter)
+    flow_integrated = zeros(flow_counter)
     flow_vertical = 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)
         flow_vertical = DiffCache(flow_vertical, chunk_sizes)
@@ -101,8 +105,13 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         edges_source,
         flow_dict,
         flow,
+        flow_prev,
+        flow_integrated,
         flow_vertical_dict,
         flow_vertical,
+        flow_vertical_prev,
+        flow_vertical_integrated,
+        config.solver.saveat,
     )
     graph = @set graph.graph_data = graph_data
 

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)
@@ -105,7 +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 = 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)
@@ -115,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, saveat_state)
     @debug "Created callbacks."
 
     # Initialize the integrator, providing all solver options as described in
@@ -131,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,
@@ -153,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
 

core/src/parameter.jl

@@ -558,8 +558,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},
+            saveat::Float64,
         },
         MetaGraphsNext.var"#11#13",
         Float64,

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

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)

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