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,46 @@ function create_callbacks(
     return callback, saved
 end
 
+"""
+Integrate flows over the last timestep
+"""
+function integrate_flows!(u, t, integrator)::Nothing
+    (; p, dt) = 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)
+
+    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 * dt
+    @. flow_vertical_integrated += flow_vertical_effective * dt
+
+    copyto!(flow_prev, flow)
+    copyto!(flow_vertical_prev, flow_vertical)
+    return nothing
+end
+
 """
 Listens for changes in condition truths.
 """
@@ -377,12 +422,34 @@ 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)
-    vcat(
-        get_tmp(integrator.p.graph[].flow_vertical, 0.0),
-        get_tmp(integrator.p.graph[].flow, 0.0),
-    )
+    (; dt, p) = integrator
+    (; graph) = p
+    (; flow_integrated, flow_vertical_integrated, saveat) = graph[]
+
+    Δt = if iszero(saveat)
+        dt
+    elseif isinf(saveat)
+        t
+    else
+        t_end = integrator.sol.prob.tspan[2]
+        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
+    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

@@ -231,20 +231,27 @@ 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,
+)::Tuple{Union{Float64, Vector{Float64}}, Vector{Float64}}
     if iszero(saveat)
         # every step
-        Float64[]
+        saveat_state = Float64[]
+        saveat_flow = Float64[]
     elseif saveat == Inf
         # only the start and end
-        [0.0, t_end]
+        saveat_state = [0.0, t_end]
+        saveat_flow = [t_end]
     elseif isfinite(saveat)
         # every saveat seconds
-        saveat
+        saveat_state = saveat
+        saveat_flow = collect(range(saveat, t_end; step = 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"

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
@@ -90,7 +92,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 +107,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 @@
     # 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_demand
     try
         parameters = Parameters(db, config)
 
@@ -79,7 +79,6 @@
         tstops_flow_boundary = 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)
-        tstops = sort(unique(vcat(tstops_flow_boundary, tstops_user_demand)))
 
         # use state
         state = load_structvector(db, config, BasinStateV1)
@@ -105,7 +104,9 @@
     @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, saveat_flow = convert_saveat(config.solver.saveat, t_end)
+    tstops = sort(unique(vcat(tstops_flow_boundary, tstops_user_demand, saveat_flow)))
     adaptive, dt = convert_dt(config.solver.dt)
 
     jac_prototype = config.solver.sparse ? get_jac_prototype(parameters) : nothing
@@ -116,7 +117,7 @@
     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
@@ -132,7 +133,7 @@
         callback,
         tstops,
         isoutofdomain = (u, p, t) -> any(<(0), u.storage),
-        saveat,
+        saveat = saveat_state,
         adaptive,
         dt,
         config.solver.dtmin,
@@ -154,17 +155,17 @@
 end
 
 "Get all saved times in seconds since start"
-timesteps(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}
-    return datetime_since.(timesteps(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(timesteps(model))
+    nsaved = length(tsaves(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

@@ -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
@@ -665,15 +663,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.(tsaves(model), config.starttime)
 
     storage = hcat([collect(u_.storage) for u_ in sol.u]...)
     level = zero(storage)

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.timesteps(model)
+    t = Ribasim.tsaves(model)
 
     p = model.integrator.p
     (; user_demand, graph, allocation, basin, level_demand) = p

core/test/config_test.jl

@@ -65,10 +65,10 @@ end
     algorithm(Solver(; algorithm = "Euler", autodiff = true))
 
     t_end = 100.0
-    @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 convert_saveat(Inf, t_end) == [0.0, t_end]
+    @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_throws ErrorException convert_saveat(-Inf, t_end)
     @test_throws ErrorException convert_saveat(NaN, t_end)