Thermo - Easily swap sa num methods, add gpu tests

.buildkite/pipeline.yml

@@ -536,6 +536,16 @@ steps:
       slurm_ntasks: 1
       slurm_gres: "gpu:1"
 
+  - label: "gpu_thermodynamics"
+    key: "gpu_thermodynamics"
+    command:
+      - "mpiexec julia --color=yes --project test/Common/Thermodynamics/runtests_gpu.jl"
+    agents:
+      config: gpu
+      queue: central
+      slurm_ntasks: 1
+      slurm_gres: "gpu:1"
+
   - label: "gpu_diagnostic_fields_test"
     key: "gpu_diagnostic_fields_test"
     command:

Manifest.toml

@@ -51,6 +51,12 @@ version = "0.1.0"
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
+[[BenchmarkTools]]
+deps = ["JSON", "Logging", "Printf", "Statistics", "UUIDs"]
+git-tree-sha1 = "9e62e66db34540a0c919d72172cc2f642ac71260"
+uuid = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+version = "0.5.0"
+
 [[CEnum]]
 git-tree-sha1 = "215a9aa4a1f23fbd05b92769fdd62559488d70e9"
 uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"

Project.toml

@@ -6,6 +6,7 @@ version = "0.3.0-DEV"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"

src/Common/Thermodynamics/Thermodynamics.jl

@@ -22,6 +22,23 @@ _molmass_ratio = molmass_ratio(param_set)
 
 Because these parameters are widely used throughout this module,
 `param_set` is an argument for many Thermodynamics functions.
+
+## Numerical methods for saturation adjustment
+
+Saturation adjustment function are designed to accept
+ - `sat_adjust_method` a type used to dispatch which numerical method to use
+
+and a function to return an instance of the numerical method. For example:
+
+ - `sa_numerical_method_ρpq` returns an instance of the numerical
+    method. One of these functions must be defined for the particular
+    numerical method and the particular formulation (`ρ-p-q_tot` in this case).
+
+The currently supported numerical methods, in RootSolvers.jl, are:
+ - `NewtonsMethod` uses Newton method with analytic gradients
+ - `NewtonsMethodAD` uses Newton method with autodiff
+ - `SecantMethod` uses Secant method
+ - `RegulaFalsiMethod` uses Regula-Falsi method
 """
 module Thermodynamics
 
@@ -55,6 +72,7 @@ print_warning() = true
 include("states.jl")
 include("relations.jl")
 include("isentropic.jl")
+include("config_numerical_method.jl")
 
 Base.broadcastable(dap::DryAdiabaticProcess) = Ref(dap)
 Base.broadcastable(phase::Phase) = Ref(phase)

src/Common/Thermodynamics/config_numerical_method.jl

@@ -0,0 +1,119 @@
+# These functions (variants of sa_numerical_method)
+# return an instance of a numerical method to solve
+# saturation adjustment, for different combinations
+# of thermodynamic variable inputs.
+
+# @print only accepts literal strings, so we must
+# branch to print which method is being used.
+function print_numerical_method(
+    ::Type{sat_adjust_method},
+) where {sat_adjust_method}
+    if sat_adjust_method <: NewtonsMethod
+        @print("    Method=NewtonsMethod")
+    elseif sat_adjust_method <: NewtonsMethodAD
+        @print("    Method=NewtonsMethodAD")
+    elseif sat_adjust_method <: SecantMethod
+        @print("    Method=SecantMethod")
+    elseif sat_adjust_method <: RegulaFalsiMethod
+        @print("    Method=RegulaFalsiMethod")
+    else
+        error("Unsupported numerical method")
+    end
+end
+
+#####
+##### Thermodynamic variable inputs: ρ, e_int, q_tot
+#####
+function sa_numerical_method(
+    ::Type{NM},
+    param_set::APS,
+    ρ::FT,
+    e_int::FT,
+    q_tot::FT,
+    phase_type::Type{<:PhaseEquil},
+) where {FT, NM <: NewtonsMethod}
+    _T_min::FT = T_min(param_set)
+    T_init =
+        max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
+    return NewtonsMethod(
+        T_init,
+        T_ -> ∂e_int_∂T(param_set, heavisided(T_), e_int, ρ, q_tot, phase_type),
+    )
+end
+
+function sa_numerical_method(
+    ::Type{NM},
+    param_set::APS,
+    ρ::FT,
+    e_int::FT,
+    q_tot::FT,
+    phase_type::Type{<:PhaseEquil},
+) where {FT, NM <: NewtonsMethodAD}
+    _T_min::FT = T_min(param_set)
+    T_init =
+        max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
+    return NewtonsMethodAD(T_init)
+end
+
+function sa_numerical_method(
+    ::Type{NM},
+    param_set::APS,
+    ρ::FT,
+    e_int::FT,
+    q_tot::FT,
+    phase_type::Type{<:PhaseEquil},
+) where {FT, NM <: SecantMethod}
+    _T_min::FT = T_min(param_set)
+    q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
+    T_2 = air_temperature(param_set, e_int, q_pt)
+    T_1 = max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
+    T_2 = bound_upper_temperature(T_1, T_2)
+    return SecantMethod(T_1, T_2)
+end
+
+function sa_numerical_method(
+    ::Type{NM},
+    param_set::APS,
+    ρ::FT,
+    e_int::FT,
+    q_tot::FT,
+    phase_type::Type{<:PhaseEquil},
+) where {FT, NM <: RegulaFalsiMethod}
+    _T_min::FT = T_min(param_set)
+    q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
+    T_2 = air_temperature(param_set, e_int, q_pt)
+    T_1 = max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
+    T_2 = bound_upper_temperature(T_1, T_2)
+    return RegulaFalsiMethod(T_1, T_2)
+end
+
+#####
+##### Thermodynamic variable inputs: ρ, p, q_tot
+#####
+
+function sa_numerical_method_ρpq(
+    ::Type{NM},
+    param_set::APS,
+    ρ::FT,
+    p::FT,
+    q_tot::FT,
+    phase_type::Type{<:PhaseEquil},
+) where {FT, NM <: NewtonsMethodAD}
+    q_pt = PhasePartition(q_tot)
+    T_init = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt)
+    return NewtonsMethodAD(T_init)
+end
+
+function sa_numerical_method_ρpq(
+    ::Type{NM},
+    param_set::APS,
+    ρ::FT,
+    p::FT,
+    q_tot::FT,
+    phase_type::Type{<:PhaseEquil},
+) where {FT, NM <: RegulaFalsiMethod}
+    q_pt = PhasePartition(q_tot)
+    T_1 = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt) - 5
+    T_2 = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt) + 5
+    return RegulaFalsiMethod(T_1, T_2)
+end