Add remaining water balance terms to basin.arrow

Technically these terms can be calculated from flow.arrow and the storage. But since calculating the water balance for a Basin is such a common operation, we include them directly. This also makes it easier to locate water balance errors in time and space, by including the error term.

At some point we may also want to add the instantaneous inflow and outflow to the Basin struct, and update them in `formulate_du!`. This would allow us to reject steps with too large of a balance error, using one of the metrics mentioned in #1271. But for now this only adds them to the output via a callback.

core/src/callback.jl

@@ -69,8 +69,8 @@ function create_callbacks(
         SavingCallback(save_vertical_flux, saved_vertical_flux; saveat, save_start = false)
     push!(callbacks, save_vertical_flux_cb)
 
-    # save the flows over time, as a Vector of the nonzeros(flow)
-    saved_flow = SavedValues(Float64, Vector{Float64})
+    # save the flows over time
+    saved_flow = SavedValues(Float64, SavedFlow)
     save_flow_cb = SavingCallback(save_flow, saved_flow; saveat, save_start = false)
     push!(callbacks, save_flow_cb)
 
@@ -138,14 +138,39 @@ end
 "Compute the average flows over the last saveat interval and write
 them to SavedValues"
 function save_flow(u, t, integrator)
-    (; flow_integrated) = integrator.p.graph[]
+    (; graph) = integrator.p
+    (; flow_integrated, flow_dict) = graph[]
+    (; node_id) = integrator.p.basin
 
     Δt = get_Δt(integrator)
     flow_mean = copy(flow_integrated)
     flow_mean ./= Δt
     fill!(flow_integrated, 0.0)
 
-    return flow_mean
+    # 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 in_id in inflow_ids(graph, basin_id)
+            q = flow_mean[flow_dict[in_id, basin_id]]
+            if q > 0
+                inflow_mean[i] += q
+            else
+                outflow_mean[i] -= q
+            end
+        end
+        for out_id in outflow_ids(graph, basin_id)
+            q = flow_mean[flow_dict[basin_id, out_id]]
+            if q > 0
+                outflow_mean[i] += q
+            else
+                inflow_mean[i] -= q
+            end
+        end
+    end
+
+    return SavedFlow(; flow = flow_mean, inflow = inflow_mean, outflow = outflow_mean)
 end
 
 "Compute the average vertical fluxes over the last saveat interval and write

core/src/model.jl

@@ -1,5 +1,5 @@
 struct SavedResults{V1 <: ComponentVector{Float64}}
-    flow::SavedValues{Float64, Vector{Float64}}
+    flow::SavedValues{Float64, SavedFlow}
     vertical_flux::SavedValues{Float64, V1}
     subgrid_level::SavedValues{Float64, Vector{Float64}}
 end

core/src/parameter.jl

@@ -134,6 +134,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
+    flow::Vector{Float64}
+    inflow::Vector{Float64}
+    outflow::Vector{Float64}
+end
+
 """
 Requirements:
 

core/src/util.jl

@@ -534,6 +534,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.

core/src/write.jl

@@ -84,31 +84,35 @@ function basin_table(
     node_id::Vector{Int32},
     storage::Vector{Float64},
     level::Vector{Float64},
+    inflow_rate::Vector{Float64},
+    outflow_rate::Vector{Float64},
+    storage_rate::Vector{Float64},
     precipitation::Vector{Float64},
     evaporation::Vector{Float64},
     drainage::Vector{Float64},
     infiltration::Vector{Float64},
+    error::Vector{Float64},
 }
     (; saved) = model
-    (; vertical_flux) = saved
-
     # 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)])
     level = vec(data.level[:, begin:(end - 1)])
+    Δstorage = vec(diff(data.storage; dims = 2))
 
     nbasin = length(data.node_id)
     ntsteps = length(data.time) - 1
     nrows = nbasin * ntsteps
 
+    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)
 
     idx_row = 0
-
-    for vec in vertical_flux.saveval
+    for vec in saved.vertical_flux.saveval
         for (precipitation_, evaporation_, drainage_, infiltration_) in
             zip(vec.precipitation, vec.evaporation, vec.drainage, vec.infiltration)
             idx_row += 1
@@ -120,17 +124,27 @@ function basin_table(
     end
 
     time = repeat(data.time[begin:(end - 1)]; inner = nbasin)
+    Δtime_seconds = 0.001 * Dates.value.(diff(data.time))
+    Δtime = repeat(Δtime_seconds; inner = nbasin)
     node_id = repeat(Int32.(data.node_id); outer = ntsteps)
+    storage_rate = Δstorage ./ Δtime
+    error =
+        inflow_rate - outflow_rate - storage_rate + precipitation - evaporation + drainage -
+        infiltration
 
     return (;
         time,
         node_id,
         storage,
         level,
+        inflow_rate,
+        outflow_rate,
+        storage_rate,
         precipitation,
         evaporation,
         drainage,
         infiltration,
+        error,
     )
 end
 
@@ -184,7 +198,7 @@ function flow_table(
     from_node_id = repeat(from_node_id; outer = ntsteps)
     to_node_type = repeat(to_node_type; outer = ntsteps)
     to_node_id = repeat(to_node_id; outer = ntsteps)
-    flow_rate = FlatVector(saveval)
+    flow_rate = FlatVector(saveval, :flow)
 
     return (;
         time,

core/test/run_models_test.jl

@@ -8,7 +8,8 @@
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/trivial/ribasim.toml")
     @test ispath(toml_path)
-    model = Ribasim.run(toml_path)
+    config = Ribasim.Config(toml_path)
+    model = Ribasim.run(config)
     @test model isa Ribasim.Model
     @test successful_retcode(model)
     (; p) = model.integrator
@@ -49,12 +50,29 @@
                 :node_id,
                 :storage,
                 :level,
+                :inflow_rate,
+                :outflow_rate,
+                :storage_rate,
                 :precipitation,
                 :evaporation,
                 :drainage,
                 :infiltration,
+                :error,
+            ),
+            (
+                DateTime,
+                Int32,
+                Float64,
+                Float64,
+                Float64,
+                Float64,
+                Float64,
+                Float64,
+                Float64,
+                Float64,
+                Float64,
+                Float64,
             ),
-            (DateTime, Int32, Float64, Float64, Float64, Float64, Float64, Float64),
         )
         @test Tables.schema(control) == Tables.Schema(
             (:time, :control_node_id, :truth_state, :control_state),
@@ -113,6 +131,14 @@
 
         @test basin.storage[1] ≈ 1.0
         @test basin.level[1] ≈ 0.044711584
+        @test basin.storage_rate[1] ≈
+              (basin.storage[2] - basin.storage[1]) / config.solver.saveat
+        @test all(==(0), basin.inflow_rate)
+        @test all(>(0), basin.outflow_rate)
+        @test flow.flow_rate[1] == basin.outflow_rate[1]
+        @test all(==(0), basin.drainage)
+        @test all(==(0), basin.infiltration)
+        @test all(q -> abs(q) < 1e-5, basin.error)
 
         # The exporter interpolates 1:1 for three subgrid elements, but shifted by 1.0 meter.
         basin_level = basin.level[1]

docs/core/usage.qmd

@@ -707,23 +707,31 @@ derivative       | Float64  | -        | -
 
 The Basin table contains:
 
-- Results of the storage and level of each basin, which are instantaneous values;
-- Results of the vertical fluxes on each basin, which are mean values over the `saveat` intervals. In the time column the start of the period is indicated.
-- The final state of the model is not part of this file, and will be placed in a separate output state file in the future.
-
-The initial condition is also written to the file.
+- Results of the storage and level of each Basin, which are instantaneous values;
+- Results of the fluxes on each Basin, which are mean values over the `saveat` intervals.
+  In the time column the start of the period is indicated.
+- The initial condition is written to the file, but the final state is not.
+  It will be placed in a separate output state file in the future.
+- The `inflow_rate` and `outflow_rate` are the sum of the flows into and out of the Basin respectively.
+  The actual flows determine in which term they are counted, not the edge direction.
+- The `storage_rate` is flow that adds to the storage in the Basin, increasing the water level.
+- The `error` is the result of adding all terms of the water balance; `inflow_rate` - `outflow_rate` - `storage_rate` + `precipitation` - `evaporation` + `drainage` - `infiltration`.
+  It can be used to check if the numerical error when solving the water balance is sufficiently small.
 
 column        | type     | unit
 ------------- | ---------| ----
 time          | DateTime | -
 node_id       | Int32    | -
 storage       | Float64  | $m^3$
 level         | Float64  | $m$
+inflow_rate   | Float64  | $m^3 s^{-1}$
+outflow_rate  | Float64  | $m^3 s^{-1}$
+storage_rate  | Float64  | $m^3 s^{-1}$
 precipitation | Float64  | $m^3 s^{-1}$
 evaporation   | Float64  | $m^3 s^{-1}$
 drainage      | Float64  | $m^3 s^{-1}$
 infiltration  | Float64  | $m^3 s^{-1}$
-
+error         | Float64  | $m^3 s^{-1}$
 
 The table is sorted by time, and per time it is sorted by `node_id`.