Save all results via saving callback (follow up) (#1896)

This is a follow up of #1825 (and a few unrelated experiments). Fixes #1837. Other changes in this PR: - Storages from the cumulative flows are computed as a sparse matrix-vector product - Storage rates are computed without involving the initial storage - Manning resistance is more stable around $\Delta h = 0$.
Deltares · Oct 10, 2024 · cf73c99 · cf73c99
1 parent 5cff21d
commit cf73c99
Show file tree

Hide file tree

Showing 13 changed files with 152 additions and 102 deletions.
diff --git a/Manifest.toml b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.10.5"
 manifest_format = "2.0"
-project_hash = "0257f2772e4bfed0b6316b6871c4495978c173b4"
+project_hash = "a2d1a982c3293971ae40e0a7bdfb40a85bd30ac1"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "eea5d80188827b35333801ef97a40c2ed653b081"
@@ -1346,7 +1346,7 @@ uuid = "295af30f-e4ad-537b-8983-00126c2a3abe"
 version = "3.6.0"
 
 [[deps.Ribasim]]
-deps = ["Accessors", "Arrow", "BasicModelInterface", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "DiffEqBase", "DiffEqCallbacks", "EnumX", "FiniteDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "LineSearches", "LinearSolve", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEqBDF", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "SparseConnectivityTracer", "StructArrays", "Tables", "TerminalLoggers", "TranscodingStreams"]
+deps = ["Accessors", "Arrow", "BasicModelInterface", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "DiffEqBase", "DiffEqCallbacks", "EnumX", "FiniteDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "LineSearches", "LinearAlgebra", "LinearSolve", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEqBDF", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "SparseConnectivityTracer", "StructArrays", "Tables", "TerminalLoggers", "TranscodingStreams"]
 path = "core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "2024.11.0"

diff --git a/core/Project.toml b/core/Project.toml
@@ -25,6 +25,7 @@ IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 Legolas = "741b9549-f6ed-4911-9fbf-4a1c0c97f0cd"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
@@ -77,6 +78,7 @@ IOCapture = "0.2"
 IterTools = "1.4"
 JuMP = "1.15"
 Legolas = "0.5"
+LinearAlgebra = "1"
 LineSearches = "7"
 LinearSolve = "2.24"
 Logging = "<0.0.1, 1"

diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl
@@ -46,6 +46,12 @@ using SciMLBase:
 # through operator overloading
 using SparseConnectivityTracer: TracerSparsityDetector, jacobian_sparsity, GradientTracer
 
+# For efficient sparse computations
+using SparseArrays: SparseMatrixCSC, spzeros
+
+# Linear algebra
+using LinearAlgebra: mul!
+
 # PreallocationTools is used because the RHS function (water_balance!) gets called with different input types
 # for u, du:
 # - Float64 for normal calls
@@ -92,7 +98,7 @@ using TerminalLoggers: TerminalLogger
 # Convenience wrapper around arrays, divides vectors in
 # separate sections which can be indexed individually.
 # Used for e.g. Basin forcing and the state vector.
-using ComponentArrays: ComponentVector, Axis
+using ComponentArrays: ComponentVector, ComponentArray, Axis, getaxes
 
 # Date and time handling; externally we use the proleptic Gregorian calendar,
 # internally we use a Float64; seconds since the start of the simulation.

diff --git a/core/src/callback.jl b/core/src/callback.jl
@@ -256,20 +256,23 @@ 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
     (; basin, water_balance_abstol, water_balance_reltol) = p
     errors = false
     current_storage = basin.current_storage[parent(u)]
-    formulate_storages!(current_storage, u, u, p, t)
+
+    # The initial storage is irrelevant for the storage rate and can only cause
+    # floating point truncation errors
+    formulate_storages!(current_storage, u, u, p, t; add_initial_storage = false)
 
     for (
         inflow_rate,
@@ -289,7 +292,7 @@ function check_water_balance_error(
         saved_flow.flow.evaporation,
         saved_flow.flow.infiltration,
         current_storage,
-        basin.storage_prev_saveat,
+        basin.Δstorage_prev_saveat,
         basin.node_id,
     )
         storage_rate = (s_now - s_prev) / Δt
@@ -304,13 +307,17 @@ 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)
         error("Too large water balance error(s) detected at t = $t")
     end
 
-    @. basin.storage_prev_saveat = current_storage
+    @. basin.Δstorage_prev_saveat = current_storage
     return nothing
 end
 

diff --git a/core/src/model.jl b/core/src/model.jl
@@ -104,6 +104,7 @@ function Model(config::Config)::Model
     du0 = zero(u0)
 
     parameters = set_state_flow_edges(parameters, u0)
+    parameters = build_flow_to_storage(parameters, u0)
     parameters = @set parameters.u_prev_saveat = zero(u0)
 
     # The Solver algorithm

diff --git a/core/src/parameter.jl b/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
 
@@ -313,7 +318,8 @@ end
     vertical_flux::V = zeros(length(node_id))
     # Initial_storage
     storage0::Vector{Float64} = zeros(length(node_id))
-    storage_prev_saveat::Vector{Float64} = zeros(length(node_id))
+    # Storage at previous saveat without storage0
+    Δ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))
@@ -840,6 +846,8 @@ const ModelGraph = MetaGraph{
     state_outflow_edge::C4 = ComponentVector()
     all_nodes_active::Base.RefValue{Bool} = Ref(false)
     tprev::Base.RefValue{Float64} = Ref(0.0)
+    # Sparse matrix for combining flows into storages
+    flow_to_storage::SparseMatrixCSC{Float64, Int64} = spzeros(1, 1)
     # Water balance tolerances
     water_balance_abstol::Float64
     water_balance_reltol::Float64

diff --git a/core/src/read.jl b/core/src/read.jl
@@ -649,7 +649,6 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
     storage0 = get_storages_from_levels(basin, state.level)
     @assert length(storage0) == n "Basin / state length differs from number of Basins"
     basin.storage0 .= storage0
-    basin.storage_prev_saveat .= storage0
     return basin
 end
 

diff --git a/core/src/solve.jl b/core/src/solve.jl
@@ -118,86 +118,40 @@ function formulate_storages!(
     du::ComponentVector,
     u::ComponentVector,
     p::Parameters,
-    t::Number,
+    t::Number;
+    add_initial_storage::Bool = true,
 )::Nothing
-    (;
-        basin,
-        flow_boundary,
-        tabulated_rating_curve,
-        pump,
-        outlet,
-        linear_resistance,
-        manning_resistance,
-        user_demand,
-        tprev,
-    ) = p
-    # Current storage: initial conditdion +
+    (; basin, flow_boundary, tprev, flow_to_storage) = p
+    # Current storage: initial condition +
     # total inflows and outflows since the start
     # of the simulation
-    current_storage .= basin.storage0
-    formulate_storage!(current_storage, basin, du, u)
+    if add_initial_storage
+        current_storage .= basin.storage0
+    else
+        current_storage .= 0.0
+    end
+    mul!(current_storage, flow_to_storage, u, 1, 1)
+    formulate_storage!(current_storage, basin, du)
     formulate_storage!(current_storage, tprev[], t, flow_boundary)
-    formulate_storage!(current_storage, t, u.tabulated_rating_curve, tabulated_rating_curve)
-    formulate_storage!(current_storage, t, u.pump, pump)
-    formulate_storage!(current_storage, t, u.outlet, outlet)
-    formulate_storage!(current_storage, t, u.linear_resistance, linear_resistance)
-    formulate_storage!(current_storage, t, u.manning_resistance, manning_resistance)
-    formulate_storage!(
-        current_storage,
-        t,
-        u.user_demand_inflow,
-        user_demand;
-        edge_volume_out = u.user_demand_outflow,
-    )
     return nothing
 end
 
 """
-The storage contributions of the forcings that are part of the state.
+The storage contributions of the forcings that are not part of the state.
 """
 function formulate_storage!(
     current_storage::AbstractVector,
     basin::Basin,
     du::ComponentVector,
-    u::ComponentVector,
 )
     (; current_cumulative_precipitation, current_cumulative_drainage) = basin
 
-    current_storage .-= u.evaporation
-    current_storage .-= u.infiltration
-
     current_cumulative_precipitation = current_cumulative_precipitation[parent(du)]
     current_cumulative_drainage = current_cumulative_drainage[parent(du)]
     current_storage .+= current_cumulative_precipitation
     current_storage .+= current_cumulative_drainage
 end
 
-"""
-Formulate storage contributions of nodes.
-"""
-function formulate_storage!(
-    current_storage::AbstractVector,
-    t::Number,
-    edge_volume_in::AbstractVector,
-    node::AbstractParameterNode;
-    edge_volume_out = nothing,
-)
-    edge_volume_out = isnothing(edge_volume_out) ? edge_volume_in : edge_volume_out
-
-    for (volume_in, volume_out, inflow_edge, outflow_edge) in
-        zip(edge_volume_in, edge_volume_out, node.inflow_edge, node.outflow_edge)
-        inflow_id = inflow_edge.edge[1]
-        if inflow_id.type == NodeType.Basin
-            current_storage[inflow_id.idx] -= volume_in
-        end
-
-        outflow_id = outflow_edge.edge[2]
-        if outflow_id.type == NodeType.Basin
-            current_storage[outflow_id.idx] += volume_out
-        end
-    end
-end
-
 """
 Formulate storage contributions of flow boundaries.
 """

diff --git a/core/src/util.jl b/core/src/util.jl
@@ -641,7 +641,7 @@ function get_variable_ref(
     variable::String;
     listen::Bool = true,
 )::Tuple{PreallocationRef, Bool}
-    (; basin, graph) = p
+    (; basin) = p
     errors = false
 
     # Only built here because it is needed to obtain indices
@@ -840,7 +840,8 @@ but the derivative is bounded at x = 0.
 """
 function relaxed_root(x, threshold)
     if abs(x) < threshold
-        1 / 4 * (x / sqrt(threshold)) * (5 - (x / threshold)^2)
+        x_scaled = x / threshold
+        sqrt(threshold) * x_scaled^3 * (9 - 5x_scaled^2) / 4
     else
         sign(x) * sqrt(abs(x))
     end
@@ -953,6 +954,56 @@ function build_state_vector(p::Parameters)
     )
 end
 
+function build_flow_to_storage(p::Parameters, u::ComponentVector)::Parameters
+    n_basins = length(p.basin.node_id)
+    n_states = length(u)
+    flow_to_storage = ComponentArray(
+        spzeros(n_basins, n_states),
+        (Axis(; basins = 1:n_basins), only(getaxes(u))),
+    )
+
+    for node_name in (
+        :tabulated_rating_curve,
+        :pump,
+        :outlet,
+        :linear_resistance,
+        :manning_resistance,
+        :user_demand,
+    )
+        node = getfield(p, node_name)
+
+        if node_name == :user_demand
+            flow_to_storage_node_inflow = view(flow_to_storage, :, :user_demand_inflow)
+            flow_to_storage_node_outflow = view(flow_to_storage, :, :user_demand_outflow)
+        else
+            flow_to_storage_node_inflow = view(flow_to_storage, :, node_name)
+            flow_to_storage_node_outflow = flow_to_storage_node_inflow
+        end
+
+        for (inflow_edge, outflow_edge) in zip(node.inflow_edge, node.outflow_edge)
+            inflow_id, node_id = inflow_edge.edge
+            if inflow_id.type == NodeType.Basin
+                flow_to_storage_node_inflow[inflow_id.idx, node_id.idx] = -1.0
+            end
+
+            outflow_id = outflow_edge.edge[2]
+            if outflow_id.type == NodeType.Basin
+                flow_to_storage_node_outflow[outflow_id.idx, node_id.idx] = 1.0
+            end
+        end
+    end
+
+    flow_to_storage_evaporation = view(flow_to_storage, :, :evaporation)
+    flow_to_storage_infiltration = view(flow_to_storage, :, :infiltration)
+
+    for i in 1:n_basins
+        flow_to_storage_evaporation[i, i] = -1.0
+        flow_to_storage_infiltration[i, i] = -1.0
+    end
+
+    @set p.flow_to_storage = parent(flow_to_storage)
+end
+
 """
 Create vectors state_inflow_edge and state_outflow_edge which give for each state
 in the state vector in order the metadata of the edge that is associated with that state.

diff --git a/core/src/write.jl b/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,