Merge #132

132: Add data collection mechanism r=charleskawczynski a=charleskawczynski

One concern we've had recently is whether our tolerance for saturation adjustment is reasonable or not because it's a bit difficult to reason about. If the tolerance is occasionally poor, we could have convergence issues with reasonable inputs and, at the moment, it's a bit difficult to know that since the inputs are not in more intuitive variables like temperature. #128 is an attempt to make sure that we start with a reasonable guess, however, knowing whether this is a good idea or not requires collecting data from a real-world run, and we don't have a clean way to do that at the moment.

This PR adds a module dedicated to collecting data, to help us better understand statistics of some important information:
 - Maximum number of iterations performed
 - Average number of iterations performed
 - Number of converged and non-converged calls (if / when we set `TD.error_on_non_convergence() = false`)

Here's a simple script for running the moist baroclinic wave in ClimaAtmos:

```julia
using Revise; include("examples/hybrid/cli_options.jl");
dict = parsed_args_per_job_id();
parsed_args = dict["sphere_baroclinic_wave_rhoe_equilmoist"];
parsed_args["enable_threading"] = false
import Thermodynamics
import RootSolvers
Thermodynamics.solution_type() = RootSolvers.VerboseSolution()
include("examples/hybrid/driver.jl")
Thermodynamics.DataCollection.print_summary()
```

At the moment, this produces

```julia
julia> Thermodynamics.DataCollection.print_summary()
┌ Info: Thermodynamics saturation_adjustment statistics:
│   max_iter = 1
│   call_counter = 15904225
│   average_max_iter = 6.287637404526156e-8
│   converged_counter = 15904225
└   non_converged_counter = 0
```
Which seems pretty good, however, this is running at coarse resolution, for only 6 days (1152 `step!`s). Runtime was about 4 min, so we can definitely crank things up.

Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

docs/src/API.md

@@ -136,3 +136,9 @@ TemperatureProfiles.DecayingTemperatureProfile
 ```@docs
 Thermodynamics.TestedProfiles
 ```
+
+## Data collection
+
+```@docs
+Thermodynamics.DataCollection
+```

src/DataCollection.jl

@@ -0,0 +1,77 @@
+"""
+    DataCollection
+
+This module is designed to help judge the accuracy and
+performance for a particular formulation, tolerance, and or
+solver configuration, by providing tools to collect various
+statistics when Thermodynamic `saturation_adjustment` is called.
+
+## Example:
+```
+import Thermodynamics as TD
+import RootSolvers as RS
+
+function do_work()
+    # Calls TD.PhaseEquil_ρeq()..., possibly many times
+end
+
+TD.solution_type() = RS.VerboseSolution()
+do_work()
+TD.DataCollection.print_summary()
+```
+!!! warn
+    This data collection was designed for unthreaded single processor
+    runs, and may not work correctly for threaded / multi-processor runs.
+"""
+module DataCollection
+
+import RootSolvers
+const RS = RootSolvers
+
+# Stats to collect
+const ref_max_iter = Ref{Int}(0)
+const ref_call_counter = Ref{Int}(0)
+const ref_converged_counter = Ref{Int}(0)
+const ref_non_converged_counter = Ref{Int}(0)
+const ref_iter_performed = Ref{Int}(0)
+
+@inline log_meta(sol::RS.CompactSolutionResults) = nothing
+
+function log_meta(sol::RS.VerboseSolutionResults)
+    ref_max_iter[] = max(ref_max_iter[], sol.iter_performed)
+    if sol.converged
+        ref_converged_counter[] += 1
+    else
+        ref_non_converged_counter[] += 1
+    end
+    ref_call_counter[] += 1
+    ref_iter_performed[] += sol.iter_performed
+    return nothing
+end
+
+function reset_stats()
+    ref_max_iter[] = 0
+    ref_call_counter[] = 0
+    ref_converged_counter[] = 0
+    ref_non_converged_counter[] = 0
+    return nothing
+end
+
+function get_data()
+    max_iter = ref_max_iter[]
+    call_counter = ref_call_counter[]
+    converged_counter = ref_converged_counter[]
+    non_converged_counter = ref_non_converged_counter[]
+    return (; max_iter, call_counter, converged_counter, non_converged_counter)
+end
+
+function print_summary(data)
+    max_iter = data.max_iter
+    call_counter = data.call_counter
+    converged_counter = data.converged_counter
+    non_converged_counter = data.non_converged_counter
+    average_max_iter = max_iter / call_counter
+    @info "Thermodynamics `saturation_adjustment` statistics:" max_iter call_counter average_max_iter converged_counter non_converged_counter
+end
+
+end # module

src/Thermodynamics.jl

@@ -74,6 +74,10 @@ print_warning() = true
 
 @inline q_pt_0(::Type{FT}) where {FT} = PhasePartition(FT(0), FT(0), FT(0))
 
+@inline solution_type() = RS.CompactSolution()
+include("DataCollection.jl")
+import .DataCollection
+
 include("states.jl")
 include("relations.jl")
 include("isentropic.jl")

src/relations.jl

@@ -1512,10 +1512,11 @@ function saturation_adjustment(
             q_tot,
             phase_type,
         ),
-        RS.CompactSolution(),
+        solution_type(),
         tol,
         maxiter,
     )
+    DataCollection.log_meta(sol)
     if !sol.converged
         if print_warning()
             KA.@print("-----------------------------------------\n")

test/relations.jl

@@ -1672,3 +1672,17 @@ end
 
 rm(joinpath(@__DIR__, "logfilepath_Float32.toml"); force = true)
 rm(joinpath(@__DIR__, "logfilepath_Float64.toml"); force = true)
+
+TD.solution_type() = RS.VerboseSolution()
+@testset "Test data collection" begin
+    ArrayType = Array{Float64}
+    FT = eltype(ArrayType)
+    param_set = parameter_set(FT)
+    profiles = TestedProfiles.PhaseEquilProfiles(param_set, ArrayType)
+    @unpack ρ, e_int, q_tot = profiles
+    ts = PhaseEquil_ρeq.(param_set, ρ, e_int, q_tot)
+    data = TD.DataCollection.get_data()
+    TD.DataCollection.print_summary(data)
+    TD.DataCollection.reset_stats()
+end
+TD.solution_type() = RS.CompactSolution()