Save balance error in saving callback

core/src/callback.jl

@@ -256,13 +256,13 @@ function save_flow(u, t, integrator)
         drainage,
         t,
     )
-    check_water_balance_error(integrator, saved_flow, Δt)
+    check_water_balance_error!(saved_flow, integrator, Δt)
     return saved_flow
 end
 
-function check_water_balance_error(
-    integrator::DEIntegrator,
+function check_water_balance_error!(
     saved_flow::SavedFlow,
+    integrator::DEIntegrator,
     Δt::Float64,
 )::Nothing
     (; u, p, t) = integrator
@@ -307,6 +307,10 @@ function check_water_balance_error(
             errors = true
             @error "Too large water balance error" id balance_error relative_error
         end
+
+        saved_flow.storage_rate[id.idx] = storage_rate
+        saved_flow.balance_error[id.idx] = balance_error
+        saved_flow.relative_error[id.idx] = relative_error
     end
     if errors
         t = datetime_since(t, p.starttime)

core/src/parameter.jl

@@ -268,6 +268,8 @@ In-memory storage of saved mean flows for writing to results.
 - `flow_boundary`: The exact integrated mean flows of flow boundaries
 - `precipitation`: The exact integrated mean precipitation
 - `drainage`: The exact integrated mean drainage
+- `balance_error`: The (absolute) water balance error
+- `relative_error`: The relative water balance error
 - `t`: Endtime of the interval over which is averaged
 """
 @kwdef struct SavedFlow{V}
@@ -277,6 +279,9 @@ In-memory storage of saved mean flows for writing to results.
     flow_boundary::Vector{Float64}
     precipitation::Vector{Float64}
     drainage::Vector{Float64}
+    storage_rate::Vector{Float64} = zero(precipitation)
+    balance_error::Vector{Float64} = zero(precipitation)
+    relative_error::Vector{Float64} = zero(precipitation)
     t::Float64
 end
 

core/src/write.jl

@@ -133,44 +133,26 @@ function basin_table(
 
     inflow_rate = FlatVector(saved.flow.saveval, :inflow)
     outflow_rate = FlatVector(saved.flow.saveval, :outflow)
-    precipitation = zeros(nrows)
-    evaporation = zeros(nrows)
-    drainage = zeros(nrows)
+    drainage = FlatVector(saved.flow.saveval, :drainage)
     infiltration = zeros(nrows)
-    balance_error = zeros(nrows)
-    relative_error = zeros(nrows)
+    evaporation = zeros(nrows)
+    precipitation = FlatVector(saved.flow.saveval, :precipitation)
+    storage_rate = FlatVector(saved.flow.saveval, :storage_rate)
+    balance_error = FlatVector(saved.flow.saveval, :balance_error)
+    relative_error = FlatVector(saved.flow.saveval, :relative_error)
 
     idx_row = 0
     for cvec in saved.flow.saveval
-        for (precipitation_, evaporation_, drainage_, infiltration_) in zip(
-            cvec.precipitation,
-            cvec.flow.evaporation,
-            cvec.drainage,
-            cvec.flow.infiltration,
-        )
+        for (evaporation_, infiltration_) in
+            zip(cvec.flow.evaporation, cvec.flow.infiltration)
             idx_row += 1
-            precipitation[idx_row] = precipitation_
             evaporation[idx_row] = evaporation_
-            drainage[idx_row] = drainage_
             infiltration[idx_row] = infiltration_
         end
     end
 
     time = repeat(data.time[begin:(end - 1)]; inner = nbasin)
-    Δtime_seconds = seconds.(diff(data.time))
-    Δtime = repeat(Δtime_seconds; inner = nbasin)
     node_id = repeat(Int32.(data.node_id); outer = ntsteps)
-    storage_rate = Δstorage ./ Δtime
-
-    for i in 1:nrows
-        total_in = inflow_rate[i] + precipitation[i] + drainage[i]
-        total_out = outflow_rate[i] + evaporation[i] + infiltration[i]
-        balance_error[i] = storage_rate[i] - (total_in - total_out)
-        mean_flow_rate = 0.5 * (total_in + total_out)
-        if mean_flow_rate != 0
-            relative_error[i] = balance_error[i] / mean_flow_rate
-        end
-    end
 
     return (;
         time,