Add remaining water balance terms to basin.arrow (#1367)

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.

---------

Co-authored-by: Bart de Koning <[email protected]>

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/read.jl

@@ -1009,13 +1009,22 @@ function exists(db::DB, tablename::String)
     return !isempty(query)
 end
 
+"""
+    seconds(period::Millisecond)::Float64
+
+Convert a period of type Millisecond to a Float64 in seconds.
+You get Millisecond objects when subtracting two DateTime objects.
+Dates.value returns the number of milliseconds.
+"""
+seconds(period::Millisecond)::Float64 = 0.001 * Dates.value(period)
+
 """
     seconds_since(t::DateTime, t0::DateTime)::Float64
 
 Convert a DateTime to a float that is the number of seconds since the start of the
 simulation. This is used to convert between the solver's inner float time, and the calendar.
 """
-seconds_since(t::DateTime, t0::DateTime)::Float64 = 0.001 * Dates.value(t - t0)
+seconds_since(t::DateTime, t0::DateTime)::Float64 = seconds(t - t0)
 
 """
     datetime_since(t::Real, t0::DateTime)::DateTime

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,33 +84,40 @@ 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},
+    balance_error::Vector{Float64},
+    relative_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)
+    balance_error = zeros(nrows)
+    relative_error = zeros(nrows)
 
     idx_row = 0
-
-    for vec in vertical_flux.saveval
+    for cvec in saved.vertical_flux.saveval
         for (precipitation_, evaporation_, drainage_, infiltration_) in
-            zip(vec.precipitation, vec.evaporation, vec.drainage, vec.infiltration)
+            zip(cvec.precipitation, cvec.evaporation, cvec.drainage, cvec.infiltration)
             idx_row += 1
             precipitation[idx_row] = precipitation_
             evaporation[idx_row] = evaporation_
@@ -120,17 +127,39 @@ function basin_table(
     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
+        storage_flow = storage_rate[i]
+        storage_increase = max(storage_flow, 0.0)
+        storage_decrease = max(-storage_flow, 0.0)
+
+        total_in = inflow_rate[i] + precipitation[i] + drainage[i] - storage_increase
+        total_out = outflow_rate[i] + evaporation[i] + infiltration[i] - storage_decrease
+        balance_error[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,
         node_id,
         storage,
         level,
+        inflow_rate,
+        outflow_rate,
+        storage_rate,
         precipitation,
         evaporation,
         drainage,
         infiltration,
+        balance_error,
+        relative_error,
     )
 end
 
@@ -184,7 +213,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,31 @@
                 :node_id,
                 :storage,
                 :level,
+                :inflow_rate,
+                :outflow_rate,
+                :storage_rate,
                 :precipitation,
                 :evaporation,
                 :drainage,
                 :infiltration,
+                :balance_error,
+                :relative_error,
+            ),
+            (
+                DateTime,
+                Int32,
+                Float64,
+                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 +133,15 @@
 
         @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-7, basin.balance_error)
+        @test all(q -> abs(q) < 0.01, basin.relative_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,33 @@ 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.
-
-column        | type     | unit
-------------- | ---------| ----
-time          | DateTime | -
-node_id       | Int32    | -
-storage       | Float64  | $m^3$
-level         | Float64  | $m$
-precipitation | Float64  | $m^3 s^{-1}$
-evaporation   | Float64  | $m^3 s^{-1}$
-drainage      | Float64  | $m^3 s^{-1}$
-infiltration  | Float64  | $m^3 s^{-1}$
-
+- 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 from other nodes 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. In the equations below this number is split out into two non-negative numbers, `storage_increase` and `storage_decrease`.
+- The `balance_error` is the difference of all Basin inflows (`total_inflow`) and outflows (`total_outflow`), that is (`inflow_rate` + `precipitation` + `drainage` - `storage_increase`) - (`outflow_rate` + `evaporation` + `infiltration` - `storage_decrease`).
+  It can be used to check if the numerical error when solving the water balance is sufficiently small.
+- The `relative_error` is the fraction of the `balance_error` over the mean of the `total_inflow` and `total_outflow`.
+
+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}$
+balance_error  | Float64  | $m^3 s^{-1}$
+relative_error | Float64  | -
 
 The table is sorted by time, and per time it is sorted by `node_id`.
 

utils/runstats.jl

@@ -106,7 +106,7 @@ toml_paths = get_testmodels()
 runs = OrderedDict{String, Any}[]
 for toml_path in toml_paths
     config = Ribasim.Config(toml_path)
-    println(basename(toml_path))
+    println(basename(dirname(toml_path)))
     # run first to compile, if this takes too long perhaps we can shorten the duration
     Ribasim.run(config)
     timed = @timed Ribasim.run(config)