Some fixes

core/src/Ribasim.jl

@@ -67,7 +67,7 @@ using SciMLBase:
     ODEFunction,
     ODEProblem,
     ODESolution,
-    VectorContinuousCallback,
+    DiscreteCallback,
     get_proposed_dt
 using SparseArrays: SparseMatrixCSC, spzeros
 using SQLite: SQLite, DB, Query, esc_id

core/src/callback.jl

@@ -21,11 +21,11 @@ function TrapezoidIntegrationCallback(integrand_func!, integral_value)::Discrete
 end
 
 function (affect!::TrapezoidIntegrationAffect)(integrator)::Nothing
-    (; integrand_func, integrand_value, integrand_value_prev, cache, integral) = affect!
-    (; dt) = integrand
+    (; integrand_func!, integrand_value, integrand_value_prev, cache, integral) = affect!
+    (; dt) = integrator
 
     copyto!(integrand_value_prev, integrand_value)
-    integrand_func(integrand_value, integrator.p)
+    integrand_func!(integrand_value, integrator.p)
 
     cache += integrand_value_prev
     cache += integrand_value
@@ -65,6 +65,11 @@ function create_callbacks(
         TrapezoidIntegrationCallback(flow_integrand!, graph[].integrated_flow)
     push!(callbacks, integrating_flow_cb)
 
+    # Save mean flows
+    saved_flow = SavedValues(Float64, SavedFlow)
+    save_flow_cb = SavingCallback(save_flow, saved_flow; saveat, save_start = false)
+    push!(callbacks, save_flow_cb)
+
     # Integrate vertical basin fluxes for BMI
     integrated_bmi_data = ComponentVector{Float64}(;
         get_tmp(basin.vertical_flux, 0)...,
@@ -124,7 +129,6 @@ function flow_integrand!(out, p)::Nothing
 
     out.flow .= flow
     vertical_flux_view(out) .= vertical_flux
-    set_inoutflows!(out, graph, basin)
     return nothing
 end
 
@@ -153,41 +157,14 @@ function allocation_integrand!(out, p)::Nothing
     return nothing
 end
 
-"Compute the average flows over the last saveat interval and write
-them to SavedValues"
 function save_flow(u, t, integrator)
-    (; graph) = integrator.p
-    (; flow_integrated, flow_dict) = graph[]
-    (; node_id) = integrator.p.basin
-
+    (; basin, graph) = integrator.p
+    (; integrated_flow) = graph[]
     Δt = get_Δt(integrator)
-    flow_mean = copy(flow_integrated)
+    flow_mean = copy(integrated_flow)
     flow_mean ./= Δt
-    fill!(flow_integrated, 0.0)
-
-    # Divide the flows over edges to Basin inflow and outflow, regardless of edge direction.
-    inflow_mean = zeros(length(node_id))
-    outflow_mean = zeros(length(node_id))
-
-    for (i, basin_id) in enumerate(node_id)
-        for inflow_id in inflow_ids(graph, basin_id)
-            q = flow_mean[flow_dict[inflow_id, basin_id]]
-            if q > 0
-                inflow_mean[i] += q
-            else
-                outflow_mean[i] -= q
-            end
-        end
-        for outflow_id in outflow_ids(graph, basin_id)
-            q = flow_mean[flow_dict[basin_id, outflow_id]]
-            if q > 0
-                outflow_mean[i] += q
-            else
-                inflow_mean[i] -= q
-            end
-        end
-    end
-
+    fill!(integrated_flow, 0.0)
+    inflow_mean, outflow_mean = compute_mean_inoutflows(flow_mean.flow, graph, basin)
     return SavedFlow(; flow = flow_mean, inflow = inflow_mean, outflow = outflow_mean)
 end
 

core/src/graph.jl

@@ -104,11 +104,19 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         flow = DiffCache(flow, chunk_sizes)
     end
 
-    flow = zero(flow_counter)
+    flow = zeros(flow_counter)
     n_basins = length(get_ids(db, "Basin"))
-    vertical_flux = zero(n_basins)
-    integrated_flow = ComponentVector{Float64}(; flow, vertical_flux...)
+    integrated_flow = ComponentVector{Float64}(;
+        flow,
+        precipitation = zeros(n_basins),
+        evaporation = zeros(n_basins),
+        drainage = zeros(n_basins),
+        infiltration = zeros(n_basins),
+    )
 
+    if config.solver.autodiff
+        flow = DiffCache(flow, chunk_sizes)
+    end
     graph_data =
         (; node_ids, edges_source, flow_dict, flow, integrated_flow, config.solver.saveat)
     graph = @set graph.graph_data = graph_data

core/src/model.jl

@@ -1,7 +1,7 @@
-struct SavedResults{V1 <: ComponentVector{Float64}, V2 <: ComponentVector{Float64}}
-    flow::IntegrandValues{Float64, V1}
+struct SavedResults{V1 <: ComponentVector{Float64}}
+    flow::SavedValues{Float64, SavedFlow}
     subgrid_level::SavedValues{Float64, Vector{Float64}}
-    integrated_bmi_data::V2
+    integrated_bmi_data::V1
 end
 
 """

core/src/parameter.jl

@@ -136,6 +136,19 @@ end
 
 abstract type AbstractParameterNode end
 
+"""
+In-memory storage of saved mean flows for writing to results.
+
+- `flow`: The mean flows on all edges
+- `inflow`: The sum of the mean flows coming into each basin
+- `outflow`: The sum of the mean flows going out of each basin
+"""
+@kwdef struct SavedFlow{V1 <: ComponentVector{Float64}}
+    flow::V1
+    inflow::Vector{Float64}
+    outflow::Vector{Float64}
+end
+
 """
 Requirements:
 

core/src/read.jl

@@ -1035,7 +1035,7 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
             integrated_flow_mapping[(node_id, node_id)] = flow_counter
         end
     end
-    integrated_flow = IntegrandValuesSum(zeros(flow_counter))
+    integrated_flow = zeros(flow_counter)
 
     return Allocation(
         Int32[],

core/src/util.jl

@@ -521,6 +521,9 @@ function Base.getindex(fv::FlatVector, i::Int)
     return v[r + 1]
 end
 
+"Construct a FlatVector from one of the fields of SavedFlow."
+FlatVector(saveval::Vector{SavedFlow}, sym::Symbol) = FlatVector(getfield.(saveval, sym))
+
 """
 Function that goes smoothly from 0 to 1 in the interval [0,threshold],
 and is constant outside this interval.
@@ -758,28 +761,34 @@ function vertical_flux_view(v::ComponentVector)
     return @view v[(:precipitation, :evaporation, :drainage, :infiltration)]
 end
 
-function set_inoutflows!(
-    v::ComponentVector{Float64},
+function compute_mean_inoutflows(
+    flow_mean::AbstractVector,
     graph::MetaGraph,
     basin::Basin,
-)::Nothing
-    for (i, basin_id) in enumerate(basin.node_id)
+)::Tuple{Vector{Float64}, Vector{Float64}}
+    (; node_id) = basin
+
+    # Divide the flows over edges to Basin inflow and outflow, regardless of edge direction.
+    inflow_mean = zeros(length(node_id))
+    outflow_mean = zeros(length(node_id))
+
+    for (i, basin_id) in enumerate(node_id)
         for inflow_edge in basin.inflow_edges[i]
             q = get_flow(graph, inflow_edge, 0)
             if q > 0
-                v.basin_inflow[i] += q
+                inflow_mean[i] += q
             else
-                v.basin_outflow[i] -= q
+                outflow_mean[i] -= q
             end
         end
         for outflow_edge in basin.outflow_edges[i]
             q = get_flow(graph, outflow_edge, 0)
             if q > 0
-                v.basin_outflow[i] += q
+                inflow_mean[i] += q
             else
-                v.basin_inflow[i] -= q
+                outflow_mean[i] -= q
             end
         end
     end
-    return nothing
+    return inflow_mean, outflow_mean
 end

core/src/write.jl

@@ -94,7 +94,7 @@ function basin_table(
     balance_error::Vector{Float64},
     relative_error::Vector{Float64},
 }
-    (; saved, integrator) = model
+    (; saved) = model
     # The last timestep is not included; there is no period over which to compute flows.
     data = get_storages_and_levels(model)
     storage = vec(data.storage[:, begin:(end - 1)])
@@ -104,17 +104,37 @@ function basin_table(
     nbasin = length(data.node_id)
     ntsteps = length(data.time) - 1
     nrows = nbasin * ntsteps
-    saveats = integrator.sol.t
-
-    inflow_rate = average_over_saveats(saved.flow, :basin_inflow, saveats)
-    outflow_rate = average_over_saveats(saved.flow, :basin_inflow, saveats)
-    precipitation = average_over_saveats(saved.flow, :precipitation, saveats)
-    evaporation = average_over_saveats(saved.flow, :evaporation, saveats)
-    drainage = average_over_saveats(saved.flow, :drainage, saveats)
-    infiltration = average_over_saveats(saved.flow, :infiltration, saveats)
+
+    inflow_rate = FlatVector(saved.flow.saveval, :inflow)
+    outflow_rate = FlatVector(saved.flow.saveval, :outflow)
+    precipitation = zeros(nrows)
+    evaporation = zeros(nrows)
+    drainage = zeros(nrows)
+    infiltration = zeros(nrows)
     balance_error = zeros(nrows)
     relative_error = zeros(nrows)
 
+    @show length(model.integrator.sol.t)
+
+    idx_row = 0
+    for cvec in saved.flow.saveval
+        for (precipitation_, evaporation_, drainage_, infiltration_) in zip(
+            cvec.flow.precipitation,
+            cvec.flow.evaporation,
+            cvec.flow.drainage,
+            cvec.flow.infiltration,
+        )
+            idx_row += 1
+            precipitation[idx_row] = precipitation_
+            evaporation[idx_row] = evaporation_
+            drainage[idx_row] = drainage_
+            infiltration[idx_row] = infiltration_
+        end
+    end
+
+    @show ntsteps
+    @show length(saved.flow.saveval)
+
     time = repeat(data.time[begin:(end - 1)]; inner = nbasin)
     Δtime_seconds = seconds.(diff(data.time))
     Δtime = repeat(Δtime_seconds; inner = nbasin)