From a58f31795755ff63ee61dcddb2af3762f660b6dc Mon Sep 17 00:00:00 2001 From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com> Date: Tue, 6 Aug 2024 11:28:35 +0200 Subject: [PATCH] Write output table with solver stats per saveat (#1677) Fixes https://github.com/Deltares/Ribasim/issues/1674. Not sure how to test this? --- core/src/callback.jl | 25 +++++++++++++++++++++++-- core/src/model.jl | 1 + core/src/parameter.jl | 8 ++++++++ core/src/write.jl | 28 ++++++++++++++++++++++++++++ core/test/run_models_test.jl | 6 ++++++ docs/reference/usage.qmd | 12 ++++++++++++ 6 files changed, 78 insertions(+), 2 deletions(-) diff --git a/core/src/callback.jl b/core/src/callback.jl index 927fddaf8..05aa766e7 100644 --- a/core/src/callback.jl +++ b/core/src/callback.jl @@ -1,4 +1,3 @@ - """ Create the different callbacks that are used to store results and feed the simulation with new data. The different callbacks @@ -44,6 +43,12 @@ function create_callbacks( save_flow_cb = SavingCallback(save_flow, saved_flow; saveat, save_start = false) push!(callbacks, save_flow_cb) + # save solver stats + saved_solver_stats = SavedValues(Float64, SolverStats) + solver_stats_cb = + SavingCallback(save_solver_stats, saved_solver_stats; saveat, save_start = true) + push!(callbacks, solver_stats_cb) + # interpolate the levels saved_subgrid_level = SavedValues(Float64, Vector{Float64}) if config.results.subgrid @@ -59,12 +64,28 @@ function create_callbacks( discrete_control_cb = FunctionCallingCallback(apply_discrete_control!) push!(callbacks, discrete_control_cb) - saved = SavedResults(saved_flow, saved_vertical_flux, saved_subgrid_level) + saved = SavedResults( + saved_flow, + saved_vertical_flux, + saved_subgrid_level, + saved_solver_stats, + ) callback = CallbackSet(callbacks...) return callback, saved end +function save_solver_stats(u, t, integrator) + (; stats) = integrator.sol + (; + time = t, + rhs_calls = stats.nf, + linear_solves = stats.nsolve, + accepted_timesteps = stats.naccept, + rejected_timesteps = stats.nreject, + ) +end + function check_negative_storage(u, t, integrator)::Nothing (; basin) = integrator.p (; node_id) = basin diff --git a/core/src/model.jl b/core/src/model.jl index 2a12beea8..c3d3b922e 100644 --- a/core/src/model.jl +++ b/core/src/model.jl @@ -2,6 +2,7 @@ struct SavedResults{V1 <: ComponentVector{Float64}} flow::SavedValues{Float64, SavedFlow} vertical_flux::SavedValues{Float64, V1} subgrid_level::SavedValues{Float64, Vector{Float64}} + solver_stats::SavedValues{Float64, SolverStats} end """ diff --git a/core/src/parameter.jl b/core/src/parameter.jl index dcecbad96..d4bfbe8c5 100644 --- a/core/src/parameter.jl +++ b/core/src/parameter.jl @@ -1,3 +1,11 @@ +const SolverStats = @NamedTuple{ + time::Float64, + rhs_calls::Int, + linear_solves::Int, + accepted_timesteps::Int, + rejected_timesteps::Int, +} + # EdgeType.flow and NodeType.FlowBoundary @enumx EdgeType flow control none @eval @enumx NodeType $(config.nodetypes...) diff --git a/core/src/write.jl b/core/src/write.jl index 5cae20042..179897883 100644 --- a/core/src/write.jl +++ b/core/src/write.jl @@ -44,6 +44,11 @@ function write_results(model::Model)::Model path = results_path(config, RESULTS_FILENAME.subgrid_level) write_arrow(path, table, compress; remove_empty_table) + # solver stats + table = solver_stats_table(model) + path = results_path(config, RESULTS_FILENAME.solver_stats) + write_arrow(path, table, compress; remove_empty_table) + @debug "Wrote results." return model end @@ -56,6 +61,7 @@ const RESULTS_FILENAME = ( allocation = "allocation.arrow", allocation_flow = "allocation_flow.arrow", subgrid_level = "subgrid_level.arrow", + solver_stats = "solver_stats.arrow", ) "Get the storage and level of all basins as matrices of nbasin × ntime" @@ -183,6 +189,28 @@ function basin_table( ) end +function solver_stats_table( + model::Model, +)::@NamedTuple{ + time::Vector{DateTime}, + rhs_calls::Vector{Int}, + linear_solves::Vector{Int}, + accepted_timesteps::Vector{Int}, + rejected_timesteps::Vector{Int}, +} + solver_stats = StructVector(model.saved.solver_stats.saveval) + (; + time = datetime_since.( + solver_stats.time[1:(end - 1)], + model.integrator.p.starttime, + ), + rhs_calls = diff(solver_stats.rhs_calls), + linear_solves = diff(solver_stats.linear_solves), + accepted_timesteps = diff(solver_stats.accepted_timesteps), + rejected_timesteps = diff(solver_stats.rejected_timesteps), + ) +end + "Create a flow result table from the saved data" function flow_table( model::Model, diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl index ca3bac404..f610930fe 100644 --- a/core/test/run_models_test.jl +++ b/core/test/run_models_test.jl @@ -29,10 +29,12 @@ flow_bytes = read(normpath(dirname(toml_path), "results/flow.arrow")) basin_bytes = read(normpath(dirname(toml_path), "results/basin.arrow")) subgrid_bytes = read(normpath(dirname(toml_path), "results/subgrid_level.arrow")) + solver_stats_bytes = read(normpath(dirname(toml_path), "results/solver_stats.arrow")) flow = Arrow.Table(flow_bytes) basin = Arrow.Table(basin_bytes) subgrid = Arrow.Table(subgrid_bytes) + solver_stats = Arrow.Table(solver_stats_bytes) @testset "Schema" begin @test Tables.schema(flow) == Tables.Schema( @@ -83,6 +85,10 @@ (:time, :subgrid_id, :subgrid_level), (DateTime, Int32, Float64), ) + @test Tables.schema(solver_stats) == Tables.Schema( + (:time, :rhs_calls, :linear_solves, :accepted_timesteps, :rejected_timesteps), + (DateTime, Int, Int, Int, Int), + ) end @testset "Results size" begin diff --git a/docs/reference/usage.qmd b/docs/reference/usage.qmd index ce18f2a64..1da163aec 100644 --- a/docs/reference/usage.qmd +++ b/docs/reference/usage.qmd @@ -316,3 +316,15 @@ column | type time | DateTime subgrid_id | Int32 subgrid_level | Float64 + +## Solver statistics - `solver_stats.arrow` + +This result file contains statistics about the solver, which can give an insight into how well the solver is performing over time. The data is solved by `saveat` (see [configuration file](#configuration-file)). `water_balance` refers to the right-hand-side function of the system of differential equations solved by the Ribasim core. + +column | type +--------------------| ----- +time | DateTime +water_balance_calls | Int +linear_solves | Int +accepted_timesteps | Int +rejected_timesteps | Int