Check water balance error at each save

core/src/callback.jl

@@ -214,6 +214,7 @@ function save_flow(u, t, integrator)
     inflow_mean = zeros(length(basin.node_id))
     outflow_mean = zeros(length(basin.node_id))
 
+    # Flow contributions from horizontal flow states
     for (flow, inflow_edge, outflow_edge) in
         zip(flow_mean, state_inflow_edge, state_outflow_edge)
         inflow_id = inflow_edge.edge[1]
@@ -235,6 +236,7 @@ function save_flow(u, t, integrator)
         end
     end
 
+    # Flow contributions from flow boundaries
     flow_boundary_mean = copy(flow_boundary.cumulative_flow_saveat) ./ Δt
     flow_boundary.cumulative_flow_saveat .= 0.0
 
@@ -253,14 +255,74 @@ function save_flow(u, t, integrator)
     @. basin.cumulative_precipitation_saveat = 0.0
     @. basin.cumulative_drainage_saveat = 0.0
 
-    return SavedFlow(;
+    saved_flow = SavedFlow(;
         flow = flow_mean,
         inflow = inflow_mean,
         outflow = outflow_mean,
         flow_boundary = flow_boundary_mean,
         precipitation,
         drainage,
     )
+    check_water_balance_error(p, saved_flow, t, Δt, u)
+    return saved_flow
+end
+
+function check_water_balance_error(
+    p::Parameters,
+    saved_flow::SavedFlow,
+    t::Float64,
+    Δt::Float64,
+    u::ComponentVector,
+)::Nothing
+    (; basin, water_balance_abstol, water_balance_reltol) = p
+    errors = false
+    current_storage = basin.current_storage[parent(u)]
+
+    for (
+        i,
+        (
+            inflow_rate,
+            outflow_rate,
+            precipitation,
+            drainage,
+            evaporation,
+            infiltration,
+            s_now,
+            s_prev,
+        ),
+    ) in enumerate(
+        zip(
+            saved_flow.inflow,
+            saved_flow.outflow,
+            saved_flow.precipitation,
+            saved_flow.drainage,
+            saved_flow.flow.evaporation,
+            saved_flow.flow.infiltration,
+            current_storage,
+            basin.storage_prev_saveat,
+        ),
+    )
+        storage_rate = (s_now - s_prev) / Δt
+        total_in = inflow_rate + precipitation + drainage
+        total_out = outflow_rate + evaporation + infiltration
+        balance_error = storage_rate - (total_in - total_out)
+        mean_flow_rate = (total_in + total_out) / 2
+        relative_error = iszero(mean_flow_rate) ? 0.0 : balance_error / mean_flow_rate
+
+        if abs(balance_error) > water_balance_abstol &&
+           abs(relative_error) > water_balance_reltol
+            errors = true
+            id = id_from_state_index(p, saved_flow.flow, i)
+            @error "Too large water balance error" id balance_error relative_error
+        end
+    end
+    if errors
+        t = datetime_since(t, p.starttime)
+        error("Too large water balance error(s) detected at t = $t")
+    end
+
+    @. basin.storage_prev_saveat = current_storage
+    return nothing
 end
 
 function save_solver_stats(u, t, integrator)

core/src/config.jl

@@ -102,6 +102,8 @@ const nodetypes = collect(keys(nodekinds))
     force_dtmin::Bool = false
     abstol::Float64 = 1e-6
     reltol::Float64 = 1e-5
+    water_balance_abstol::Float64 = 1e-6
+    water_balance_reltol::Float64 = 1e-6
     maxiters::Int = 1e9
     sparse::Bool = true
     autodiff::Bool = true

core/src/parameter.jl

@@ -227,6 +227,8 @@ edge: (from node ID, to node ID)
     edge::Tuple{NodeID, NodeID}
 end
 
+Base.length(::EdgeMetadata) = 1
+
 """
 The update of an parameter given by a value and a reference to the target
 location of the variable in memory
@@ -307,6 +309,7 @@ end
     vertical_flux_bmi::V2 = zeros(length(node_id))
     # Initial_storage
     storage0::Vector{Float64} = zeros(length(node_id))
+    storage_prev_saveat::Vector{Float64} = zeros(length(node_id))
     # Analytically integrated forcings
     cumulative_precipitation::Vector{Float64} = zeros(length(node_id))
     cumulative_drainage::Vector{Float64} = zeros(length(node_id))
@@ -810,7 +813,7 @@ const ModelGraph = MetaGraph{
     Float64,
 }
 
-@kwdef struct Parameters{C1, C2, V1, V2}
+@kwdef struct Parameters{C1, C2, C3, C4, V1, V2}
     starttime::DateTime
     graph::ModelGraph
     allocation::Allocation
@@ -831,8 +834,11 @@ const ModelGraph = MetaGraph{
     flow_demand::FlowDemand
     subgrid::Subgrid
     # Per state the in- and outflow edges associated with that state (if theu exist)
-    state_inflow_edge::Vector{EdgeMetadata} = EdgeMetadata[]
-    state_outflow_edge::Vector{EdgeMetadata} = EdgeMetadata[]
+    state_inflow_edge::C3 = ComponentVector()
+    state_outflow_edge::C4 = ComponentVector()
     all_nodes_active::Base.RefValue{Bool} = Ref(false)
     tprev::Base.RefValue{Float64} = Ref(0.0)
+    # Water balance tolerances
+    water_balance_abstol::Float64
+    water_balance_reltol::Float64
 end

core/src/read.jl

@@ -654,6 +654,7 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
     )
 
     basin = @set basin.storage0 = get_storages_from_levels(basin, state.level)
+    basin = @set basin.storage_prev_saveat = copy(basin.storage0)
     @assert length(basin.storage0) == n "Basin / state length differs from number of Basins"
     return basin
 end
@@ -1289,6 +1290,8 @@ function Parameters(db::DB, config::Config)::Parameters
         level_demand,
         flow_demand,
         subgrid,
+        config.solver.water_balance_abstol,
+        config.solver.water_balance_reltol,
     )
 
     collect_control_mappings!(p)

core/src/util.jl

@@ -954,8 +954,9 @@ end
 function set_state_flow_edges(p::Parameters, u0::ComponentVector)::Parameters
     (; user_demand, graph) = p
 
-    state_inflow_edge = EdgeMetadata[]
-    state_outflow_edge = EdgeMetadata[]
+    components = Symbol[]
+    state_inflow_edges = Vector{EdgeMetadata}[]
+    state_outflow_edges = Vector{EdgeMetadata}[]
 
     placeholder_edge = EdgeMetadata(
         0,
@@ -968,6 +969,9 @@ function set_state_flow_edges(p::Parameters, u0::ComponentVector)::Parameters
     for node_name in keys(u0)
         if hasfield(Parameters, node_name)
             node::AbstractParameterNode = getfield(p, node_name)
+            push!(components, node_name)
+            state_inflow_edges_component = EdgeMetadata[]
+            state_outflow_edges_component = EdgeMetadata[]
             for id in node.node_id
                 inflow_ids_ = collect(inflow_ids(p.graph, id))
                 outflow_ids_ = collect(outflow_ids(p.graph, id))
@@ -980,7 +984,7 @@ function set_state_flow_edges(p::Parameters, u0::ComponentVector)::Parameters
                 else
                     error()
                 end
-                push!(state_inflow_edge, inflow_edge)
+                push!(state_inflow_edges_component, inflow_edge)
 
                 outflow_edge = if length(outflow_ids_) == 0
                     placeholder_edge
@@ -990,20 +994,26 @@ function set_state_flow_edges(p::Parameters, u0::ComponentVector)::Parameters
                 else
                     error()
                 end
-                push!(state_outflow_edge, outflow_edge)
+                push!(state_outflow_edges_component, outflow_edge)
             end
+            push!(state_inflow_edges, state_inflow_edges_component)
+            push!(state_outflow_edges, state_outflow_edges_component)
         elseif startswith(String(node_name), "user_demand")
+            push!(components, node_name)
             placeholder_edges = fill(placeholder_edge, length(user_demand.node_id))
             if node_name == :user_demand_inflow
-                append!(state_inflow_edge, user_demand.inflow_edge)
-                append!(state_outflow_edge, placeholder_edges)
+                push!(state_inflow_edges, user_demand.inflow_edge)
+                push!(state_outflow_edges, placeholder_edges)
             elseif node_name == :user_demand_outflow
-                append!(state_inflow_edge, placeholder_edges)
-                append!(state_outflow_edge, user_demand.outflow_edge)
+                push!(state_inflow_edges, placeholder_edges)
+                push!(state_outflow_edges, user_demand.outflow_edge)
             end
         end
     end
 
+    state_inflow_edge = ComponentVector(NamedTuple(zip(components, state_inflow_edges)))
+    state_outflow_edge = ComponentVector(NamedTuple(zip(components, state_outflow_edges)))
+
     p = @set p.state_inflow_edge = state_inflow_edge
     p = @set p.state_outflow_edge = state_outflow_edge
     return p

core/test/docs.toml

@@ -33,6 +33,8 @@ dtmax = 0.0         # optional, default length of simulation
 force_dtmin = false # optional, default false
 abstol = 1e-6       # optional, default 1e-6
 reltol = 1e-5       # optional, default 1e-5
+water_balance_abstol = 1e-6 # optional, default 1e-6
+water_balance_reltol = 1e-6 # optional, default 1e-6
 maxiters = 1e9      # optional, default 1e9
 sparse = true       # optional, default true
 autodiff = true     # optional, default true