diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 9cb7caaa610..5fdbbbbdaca 100644 --- a/.buildkite/pipeline.yml +++ b/.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: diff --git a/Manifest.toml b/Manifest.toml index 218be65126b..b4c212b9d77 100644 --- a/Manifest.toml +++ b/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" diff --git a/Project.toml b/Project.toml index 8c14f1b2ed0..db6c6ae3b51 100644 --- a/Project.toml +++ b/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" diff --git a/src/Common/Thermodynamics/Thermodynamics.jl b/src/Common/Thermodynamics/Thermodynamics.jl index cf432a7f4b0..0a6b4f3ae18 100644 --- a/src/Common/Thermodynamics/Thermodynamics.jl +++ b/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) diff --git a/src/Common/Thermodynamics/config_numerical_method.jl b/src/Common/Thermodynamics/config_numerical_method.jl new file mode 100644 index 00000000000..bd6e502c8ab --- /dev/null +++ b/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 diff --git a/src/Common/Thermodynamics/relations.jl b/src/Common/Thermodynamics/relations.jl index ad88ca0616b..61cd5bfb2bb 100644 --- a/src/Common/Thermodynamics/relations.jl +++ b/src/Common/Thermodynamics/relations.jl @@ -1260,6 +1260,7 @@ end """ saturation_adjustment( + sat_adjust_method, param_set, e_int, ρ, @@ -1271,6 +1272,8 @@ end Compute the temperature that is consistent with + - `sat_adjust_method` the numerical method to use. + See the [`Thermodynamics`](@ref) for options. - `param_set` an `AbstractParameterSet`, see the [`Thermodynamics`](@ref) for more details - `e_int` internal energy - `ρ` (moist-)air density @@ -1288,6 +1291,7 @@ using Newtons method with analytic gradients. See also [`saturation_adjustment`](@ref). """ function saturation_adjustment( + ::Type{sat_adjust_method}, param_set::APS, e_int::FT, ρ::FT, @@ -1295,7 +1299,7 @@ function saturation_adjustment( phase_type::Type{<:PhaseEquil}, maxiter::Int, temperature_tol::FT, -) where {FT <: Real} +) where {FT <: Real, sat_adjust_method} _T_min::FT = T_min(param_set) _cv_d = FT(cv_d(param_set)) # Convert temperature tolerance to a convergence criterion on internal energy residuals @@ -1330,20 +1334,17 @@ function saturation_adjustment( internal_energy_sat( param_set, heavisided(T), - ρ, - q_tot, + oftype(T, ρ), + oftype(T, q_tot), phase_type, ) - e_int, - NewtonsMethod( - T_1, - T_ -> ∂e_int_∂T( - param_set, - heavisided(T_), - e_int, - ρ, - q_tot, - phase_type, - ), + sa_numerical_method( + sat_adjust_method, + param_set, + ρ, + e_int, + q_tot, + phase_type, ), CompactSolution(), tol, @@ -1353,21 +1354,13 @@ function saturation_adjustment( if print_warning() @print("-----------------------------------------\n") @print("maxiter reached in saturation_adjustment:\n") - @print( - " e_int=", - e_int, - ", ρ=", - ρ, - ", q_tot=", - q_tot, - ", T = ", - sol.root, - ", maxiter=", - maxiter, - ", tol=", - tol.tol, - "\n" - ) + print_numerical_method(sat_adjust_method) + @print(", e_int=", e_int) + @print(", ρ=", ρ) + @print(", q_tot=", q_tot) + @print(", T=", sol.root) + @print(", maxiter=", maxiter) + @print(", tol=", tol.tol, "\n") end if error_on_non_convergence() error("Failed to converge with printed set of inputs.") @@ -1380,6 +1373,7 @@ end """ saturation_adjustment_ρpq( + sat_adjust_method, param_set, ρ, p, @@ -1389,6 +1383,8 @@ end temperature_tol ) Compute the temperature that is consistent with + - `sat_adjust_method` the numerical method to use. + See the [`Thermodynamics`](@ref) for options. - `param_set` an `AbstractParameterSet`, see the [`Thermodynamics`](@ref) for more details - `ρ` (moist-)air density - `p` pressure @@ -1410,6 +1406,7 @@ using Newtons method using ForwardDiff. See also [`saturation_adjustment`](@ref). """ function saturation_adjustment_ρpq( + ::Type{sat_adjust_method}, param_set::APS, ρ::FT, p::FT, @@ -1417,10 +1414,8 @@ function saturation_adjustment_ρpq( phase_type::Type{<:PhaseEquil}, maxiter::Int, temperature_tol::FT = eps(FT), -) where {FT <: Real} +) where {FT <: Real, sat_adjust_method} tol = SolutionTolerance(temperature_tol) - q_pt = PhasePartition(q_tot) - T_init = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt) # Use `oftype` to preserve diagonalized type signatures: sol = find_zero( T -> @@ -1436,7 +1431,14 @@ function saturation_adjustment_ρpq( phase_type, ), ), - NewtonsMethodAD(T_init), + sa_numerical_method_ρpq( + sat_adjust_method, + param_set, + ρ, + p, + q_tot, + phase_type, + ), CompactSolution(), tol, maxiter, @@ -1445,21 +1447,13 @@ function saturation_adjustment_ρpq( if print_warning() @print("-----------------------------------------\n") @print("maxiter reached in saturation_adjustment_ρpq:\n") - @print( - " ρ=", - ρ, - ", p=", - p, - ", q_tot=", - q_tot, - ", T = ", - sol.root, - ", maxiter=", - maxiter, - ", tol=", - tol.tol, - "\n" - ) + print_numerical_method(sat_adjust_method) + @print(", ρ=", ρ) + @print(", p=", p) + @print(", q_tot=", q_tot) + @print(", T=", sol.root) + @print(", maxiter=", maxiter) + @print(", tol=", tol.tol, "\n") end if error_on_non_convergence() error("Failed to converge with printed set of inputs.") @@ -1493,122 +1487,6 @@ saturation adjustment using Secant method return min(T_1 + ΔT_max(FT), T_2) end -""" - saturation_adjustment_SecantMethod( - param_set, - e_int, - ρ, - q_tot, - phase_type, - maxiter, - temperature_tol - ) - -Compute the temperature `T` that is consistent with - - - `param_set` an `AbstractParameterSet`, see the [`Thermodynamics`](@ref) for more details - - `e_int` internal energy - - `ρ` (moist-)air density - - `q_tot` total specific humidity - - `phase_type` a thermodynamic state type - - `maxiter` maximum iterations for non-linear equation solve - - `temperature_tol` temperature tolerance - -by finding the root of - -`e_int - internal_energy_sat(param_set, T, ρ, q_tot, phase_type) = 0` - -See also [`saturation_adjustment_given_ρθq`](@ref). -""" -function saturation_adjustment_SecantMethod( - param_set::APS, - e_int::FT, - ρ::FT, - q_tot::FT, - phase_type::Type{<:PhaseEquil}, - maxiter::Int, - temperature_tol::FT, -) where {FT <: Real} - _T_min::FT = T_min(param_set) - _cv_d = FT(cv_d(param_set)) - # Convert temperature tolerance to a convergence criterion on internal energy residuals - tol = ResidualTolerance(temperature_tol * _cv_d) - - T_1 = max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor - q_v_sat = q_vap_saturation(param_set, T_1, ρ, phase_type) - unsaturated = q_tot <= q_v_sat - if unsaturated && T_1 > _T_min - return T_1 - else - - _T_freeze::FT = T_freeze(param_set) - e_int_upper = internal_energy_sat( - param_set, - _T_freeze + temperature_tol / 2, # /2 => resulting interval is `temperature_tol` wide - ρ, - q_tot, - phase_type, - ) - e_int_lower = internal_energy_sat( - param_set, - _T_freeze - temperature_tol / 2, # /2 => resulting interval is `temperature_tol` wide - ρ, - q_tot, - phase_type, - ) - if e_int_lower < e_int < e_int_upper - return _T_freeze - else - # FIXME here: need to revisit bounds for saturation adjustment to guarantee bracketing of zero. - T_2 = air_temperature( - param_set, - e_int, - PhasePartition(q_tot, FT(0), q_tot), - ) # Assume all ice - T_2 = bound_upper_temperature(T_1, T_2) - sol = find_zero( - T -> - internal_energy_sat( - param_set, - heavisided(T), - ρ, - q_tot, - phase_type, - ) - e_int, - SecantMethod(T_1, T_2), - CompactSolution(), - tol, - maxiter, - ) - if !sol.converged - if print_warning() - @print("-----------------------------------------\n") - @print("maxiter reached in saturation_adjustment_SecantMethod:\n") - @print( - " e_int=", - e_int, - ", ρ=", - ρ, - ", q_tot=", - q_tot, - ", T = ", - sol.root, - ", maxiter=", - maxiter, - ", tol=", - tol.tol, - "\n" - ) - end - if error_on_non_convergence() - error("Failed to converge with printed set of inputs.") - end - end - return sol.root - end - end -end - """ saturation_adjustment_given_ρθq( param_set, @@ -1687,21 +1565,13 @@ function saturation_adjustment_given_ρθq( if print_warning() @print("-----------------------------------------\n") @print("maxiter reached in saturation_adjustment_given_ρθq:\n") - @print( - " ρ=", - ρ, - ", θ_liq_ice=", - θ_liq_ice, - ", q_tot=", - q_tot, - ", T = ", - sol.root, - ", maxiter=", - maxiter, - ", tol=", - tol.tol, - "\n" - ) + @print(" Method=SecantMethod") + @print(", ρ=", ρ) + @print(", θ_liq_ice=", θ_liq_ice) + @print(", q_tot=", q_tot) + @print(", T=", sol.root) + @print(", maxiter=", maxiter) + @print(", tol=", tol.tol, "\n") end if error_on_non_convergence() error("Failed to converge with printed set of inputs.") @@ -1796,21 +1666,13 @@ function saturation_adjustment_given_pθq( if print_warning() @print("-----------------------------------------\n") @print("maxiter reached in saturation_adjustment_given_pθq:\n") - @print( - " p=", - p, - ", θ_liq_ice=", - θ_liq_ice, - ", q_tot=", - q_tot, - ", T = ", - sol.root, - ", maxiter=", - maxiter, - ", tol=", - tol.tol, - "\n" - ) + @print(" Method=SecantMethod") + @print(", p=", p) + @print(", θ_liq_ice=", θ_liq_ice) + @print(", q_tot=", q_tot) + @print(", T=", sol.root) + @print(", maxiter=", maxiter) + @print(", tol=", tol.tol, "\n") end if error_on_non_convergence() error("Failed to converge with printed set of inputs.") @@ -1998,21 +1860,13 @@ function temperature_and_humidity_given_TᵥρRH( if print_warning() @print("-----------------------------------------\n") @print("maxiter reached in temperature_and_humidity_given_TᵥρRH:\n") - @print( - " T_virt=", - T_virt, - ", RH=", - RH, - ", ρ=", - ρ, - ", T = ", - sol.root, - ", maxiter=", - maxiter, - ", tol=", - tol.tol, - "\n" - ) + @print(" Method=SecantMethod") + @print(", T_virt=", T_virt) + @print(", RH=", RH) + @print(", ρ=", ρ) + @print(", T=", sol.root) + @print(", maxiter=", maxiter) + @print(", tol=", tol.tol, "\n") end if error_on_non_convergence() error("Failed to converge with printed set of inputs.") @@ -2104,25 +1958,15 @@ function air_temperature_given_θρq_nonlinear( if print_warning() @print("-----------------------------------------\n") @print("maxiter reached in air_temperature_given_θρq_nonlinear:\n") - @print( - " θ_liq_ice=", - θ_liq_ice, - ", ρ=", - ρ, - ", q.tot=", - q.tot, - ", q.liq = ", - q.liq, - ", q.ice = ", - q.ice, - ", T = ", - sol.root, - ", maxiter=", - maxiter, - ", tol=", - tol.tol, - "\n" - ) + @print(" Method=SecantMethod") + @print(", θ_liq_ice=", θ_liq_ice) + @print(", ρ=", ρ) + @print(", q.tot=", q.tot) + @print(", q.liq=", q.liq) + @print(", q.ice=", q.ice) + @print(", T=", sol.root) + @print(", maxiter=", maxiter) + @print(", tol=", tol.tol, "\n") end if error_on_non_convergence() error("Failed to converge with printed set of inputs.") diff --git a/src/Common/Thermodynamics/states.jl b/src/Common/Thermodynamics/states.jl index a56b3e9083a..5c00dc02148 100644 --- a/src/Common/Thermodynamics/states.jl +++ b/src/Common/Thermodynamics/states.jl @@ -201,6 +201,7 @@ end """ PhaseEquil + Moist thermodynamic phase, given - `param_set` an `AbstractParameterSet`, see the [`Thermodynamics`](@ref) for more details - `e_int` internal energy @@ -209,9 +210,8 @@ Moist thermodynamic phase, given and, optionally - `maxiter` maximum iterations for saturation adjustment - `temperature_tol` temperature tolerance for saturation adjustment - - `sat_adjust` function pointer to particular saturation adjustment method, options include - - `saturation_adjustment` uses Newtons method with analytic gradients - - `saturation_adjustment_SecantMethod` uses Secant method + - `sat_adjust_method` the numerical method to use. + See the [`Thermodynamics`](@ref) for options. """ function PhaseEquil( param_set::APS, @@ -220,11 +220,12 @@ function PhaseEquil( q_tot::FT, maxiter::Int = 8, temperature_tol::FT = FT(1e-1), - sat_adjust::Function = saturation_adjustment, -) where {FT <: Real} + ::Type{sat_adjust_method} = NewtonsMethod, +) where {FT <: Real, sat_adjust_method} phase_type = PhaseEquil q_tot_safe = clamp(q_tot, FT(0), FT(1)) - T = sat_adjust( + T = saturation_adjustment( + sat_adjust_method, param_set, e_int, ρ, @@ -236,6 +237,30 @@ function PhaseEquil( return PhaseEquil{FT, typeof(param_set)}(param_set, e_int, ρ, q_tot_safe, T) end +# Convenience method for comparing Numerical +# methods without having to specify maxiter +# and temperature_tol. maxiter and temperature_tol +# should be in sync with the PhaseEquil(...) constructor +function PhaseEquil_dev_only( + param_set::APS, + e_int::FT, + ρ::FT, + q_tot::FT; + maxiter::Int = 8, + temperature_tol::FT = FT(1e-1), + sat_adjust_method::Type{NM} = NewtonsMethod, +) where {FT <: Real, NM} + return PhaseEquil( + param_set, + e_int, + ρ, + q_tot, + maxiter, + temperature_tol, + sat_adjust_method, + ) +end + """ PhaseEquil_ρθq(param_set, ρ, θ_liq_ice, q_tot) @@ -364,6 +389,9 @@ Constructs a [`PhaseEquil`](@ref) thermodynamic state from temperature. - `p` pressure - `q_tot` total specific humidity - `perform_sat_adjust` Bool indicating to perform saturation adjustment + - `maxiter` maximum number of iterations to perform in saturation adjustment + - `sat_adjust_method` the numerical method to use. + See the [`Thermodynamics`](@ref) for options. """ function PhaseEquil_ρpq( param_set::APS, @@ -372,11 +400,13 @@ function PhaseEquil_ρpq( q_tot::FT, perform_sat_adjust = false, maxiter::Int = 5, -) where {FT <: Real} + ::Type{sat_adjust_method} = NewtonsMethodAD, +) where {FT <: Real, sat_adjust_method} phase_type = PhaseEquil if perform_sat_adjust T = saturation_adjustment_ρpq( + sat_adjust_method, param_set, ρ, p, diff --git a/test/Common/Thermodynamics/profiles.jl b/test/Common/Thermodynamics/profiles.jl index 2c1718e0668..37b0e3bfb86 100644 --- a/test/Common/Thermodynamics/profiles.jl +++ b/test/Common/Thermodynamics/profiles.jl @@ -61,6 +61,30 @@ end Base.IteratorSize(::ProfileSet) = Base.HasLength() Base.length(ps::ProfileSet) = length(ps.z) +# Since we use `rand` to generate the ProfileSet, +# just initialize on the CPU, and provide convert +# function to move arrays to the GPU. +convert_profile_set(ps::ProfileSet, ArrayType, slice) = ProfileSet( + ArrayType(ps.z[slice]), + ArrayType(ps.T[slice]), + ArrayType(ps.p[slice]), + ArrayType(ps.RS[slice]), + ArrayType(ps.e_int[slice]), + ArrayType(ps.ρ[slice]), + ArrayType(ps.θ_liq_ice[slice]), + ArrayType(ps.q_tot[slice]), + ArrayType(ps.q_liq[slice]), + ArrayType(ps.q_ice[slice]), + PhasePartition.(ps.q_tot[slice], ps.q_liq[slice], ps.q_ice[slice]), + ArrayType(ps.RH[slice]), + ArrayType(ps.e_pot[slice]), + ArrayType(ps.u[slice]), + ArrayType(ps.v[slice]), + ArrayType(ps.w[slice]), + ArrayType(ps.e_kin[slice]), + ps.phase_type, +) + """ input_config( ArrayType; @@ -85,7 +109,7 @@ function input_config( z_range = ArrayType(range(0, stop = 2.5e4, length = n)) relative_sat1 = ArrayType(range(0, stop = 1, length = n_RS1)) relative_sat2 = ArrayType(range(1, stop = 1.02, length = n_RS2)) - relative_sat = [relative_sat1..., relative_sat2...] + relative_sat = vcat(relative_sat1, relative_sat2) return z_range, relative_sat, T_surface, T_min end @@ -166,12 +190,12 @@ function PhaseDryProfiles( # Additional variables q_tot = similar(T) fill!(q_tot, 0) - q_pt = PhasePartition_equil.(Ref(param_set), T, ρ, q_tot, Ref(phase_type)) - e_int = internal_energy.(Ref(param_set), T, q_pt) - θ_liq_ice = liquid_ice_pottemp.(Ref(param_set), T, ρ, q_pt) + q_pt = PhasePartition_equil.(param_set, T, ρ, q_tot, phase_type) + e_int = internal_energy.(param_set, T, q_pt) + θ_liq_ice = liquid_ice_pottemp.(param_set, T, ρ, q_pt) q_liq = getproperty.(q_pt, :liq) q_ice = getproperty.(q_pt, :ice) - RH = relative_humidity.(Ref(param_set), T, p, Ref(phase_type), q_pt) + RH = relative_humidity.(param_set, T, p, phase_type, q_pt) e_pot = _grav * z Random.seed!(15) u = rand(FT, size(T)) * 50 diff --git a/test/Common/Thermodynamics/runtests.jl b/test/Common/Thermodynamics/runtests.jl index 711fae96835..90c55842e17 100644 --- a/test/Common/Thermodynamics/runtests.jl +++ b/test/Common/Thermodynamics/runtests.jl @@ -2,6 +2,7 @@ using Test using ClimateMachine.Thermodynamics using ClimateMachine.TemperatureProfiles using UnPack +using BenchmarkTools using NCDatasets using Random using RootSolvers @@ -321,7 +322,8 @@ end q_tot = FT(0) ρ = FT(1) phase_type = PhaseEquil - @test TD.saturation_adjustment_SecantMethod( + @test TD.saturation_adjustment( + SecantMethod, param_set, internal_energy_sat(param_set, 300.0, ρ, q_tot, phase_type), ρ, @@ -332,6 +334,7 @@ end ) ≈ 300.0 @test abs( TD.saturation_adjustment( + NewtonsMethod, param_set, internal_energy_sat(param_set, 300.0, ρ, q_tot, phase_type), ρ, @@ -345,7 +348,8 @@ end q_tot = FT(0.21) ρ = FT(0.1) @test isapprox( - TD.saturation_adjustment_SecantMethod( + TD.saturation_adjustment( + SecantMethod, param_set, internal_energy_sat(param_set, 200.0, ρ, q_tot, phase_type), ρ, @@ -359,6 +363,7 @@ end ) @test abs( TD.saturation_adjustment( + NewtonsMethod, param_set, internal_energy_sat(param_set, 200.0, ρ, q_tot, phase_type), ρ, @@ -468,16 +473,7 @@ end @test all(air_temperature.(ts) .== Ref(_T_freeze)) # Args needs to be in sync with PhaseEquil: - ts = - PhaseEquil.( - param_set, - _e_int, - ρ, - q_tot, - 8, - FT(1e-1), - TD.saturation_adjustment_SecantMethod, - ) + ts = PhaseEquil.(param_set, _e_int, ρ, q_tot, 8, FT(1e-1), SecantMethod) @test all(air_temperature.(ts) .== Ref(_T_freeze)) # PhaseEquil @@ -528,25 +524,8 @@ end # PhaseEquil ts_exact = - PhaseEquil.( - param_set, - e_int, - ρ, - q_tot, - 100, - FT(1e-3), - TD.saturation_adjustment_SecantMethod, - ) - ts = - PhaseEquil.( - param_set, - e_int, - ρ, - q_tot, - 35, - FT(1e-1), - TD.saturation_adjustment_SecantMethod, - ) # Needs to be in sync with default + PhaseEquil.(param_set, e_int, ρ, q_tot, 100, FT(1e-3), SecantMethod) + ts = PhaseEquil.(param_set, e_int, ρ, q_tot, 35, FT(1e-1), SecantMethod) # Needs to be in sync with default # Should be machine accurate (because ts contains `e_int`,`ρ`,`q_tot`): @test all(compare_moisture.(ts, ts_exact)) @test all(internal_energy.(ts) .≈ internal_energy.(ts_exact)) @@ -650,6 +629,7 @@ end @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles @test_throws ErrorException TD.saturation_adjustment.( + NewtonsMethod, param_set, e_int, ρ, @@ -659,7 +639,8 @@ end FT(1e-10), ) - @test_throws ErrorException TD.saturation_adjustment_SecantMethod.( + @test_throws ErrorException TD.saturation_adjustment.( + SecantMethod, param_set, e_int, ρ, @@ -710,6 +691,7 @@ end ) @test_throws ErrorException TD.saturation_adjustment_ρpq.( + NewtonsMethodAD, param_set, ρ, p, @@ -769,16 +751,7 @@ end @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles # PhaseEquil - ts = - PhaseEquil.( - param_set, - e_int, - ρ, - q_tot, - 40, - FT(1e-1), - Ref(TD.saturation_adjustment_SecantMethod), - ) + ts = PhaseEquil.(param_set, e_int, ρ, q_tot, 40, FT(1e-1), SecantMethod) @test all(internal_energy.(ts) .≈ e_int) @test all(getproperty.(PhasePartition.(ts), :tot) .≈ q_tot) @test all(air_density.(ts) .≈ ρ) @@ -1182,3 +1155,48 @@ end @test all(getproperty.(q_pt, :tot) .≈ (nt.q_pt.tot for nt in profiles)) @test all(phase_type .== (nt.phase_type for nt in profiles)) end + +@testset "Thermodynamics - Performance" begin + ArrayType = Array{Float64} + FT = eltype(ArrayType) + profiles = PhaseEquilProfiles(param_set, ArrayType) + + @unpack e_int, ρ, q_tot = profiles + + @btime TD.PhaseEquil_dev_only.( + $param_set, + $e_int, + $ρ, + $q_tot; + sat_adjust_method = NewtonsMethod, + ) + + @btime TD.PhaseEquil_dev_only.( + $param_set, + $e_int, + $ρ, + $q_tot; + sat_adjust_method = RegulaFalsiMethod, + maxiter = 20, + ) + + # Fails to converge: + # @btime TD.PhaseEquil_dev_only.( + # $param_set, + # $e_int, + # $ρ, + # $q_tot; + # sat_adjust_method = NewtonsMethodAD, + # maxiter = 50, + # ) + + @btime TD.PhaseEquil_dev_only.( + $param_set, + $e_int, + $ρ, + $q_tot; + sat_adjust_method = SecantMethod, + maxiter = 50, + ) + +end diff --git a/test/Common/Thermodynamics/runtests_gpu.jl b/test/Common/Thermodynamics/runtests_gpu.jl new file mode 100644 index 00000000000..05948884cab --- /dev/null +++ b/test/Common/Thermodynamics/runtests_gpu.jl @@ -0,0 +1,92 @@ +using Test +using KernelAbstractions +using ClimateMachine.Thermodynamics +using ClimateMachine.TemperatureProfiles +using UnPack +using CUDA +using Random +using RootSolvers +const TD = Thermodynamics + +using LinearAlgebra +using CLIMAParameters +using CLIMAParameters.Planet + +struct EarthParameterSet <: AbstractEarthParameterSet end +const param_set = EarthParameterSet() + +using ClimateMachine +ClimateMachine.init() +ArrayType = ClimateMachine.array_type() + +@show ArrayType + +device(::Type{T}) where {T <: Array} = CPU() +device(::Type{T}) where {T <: CuArray} = CUDADevice() + +include("profiles.jl") + +@kernel function test_thermo_kernel!( + param_set, + dst::AbstractArray{FT}, + e_int, + ρ, + p, + q_tot, +) where {FT} + i = @index(Group, Linear) + @inbounds begin + + ts = PhaseEquil(param_set, FT(e_int[i]), FT(ρ[i]), FT(q_tot[i])) + dst[1, i] = air_temperature(ts) + + ts_ρpq = PhaseEquil_ρpq( + param_set, + FT(ρ[i]), + FT(p[i]), + FT(q_tot[i]), + true, + 100, + RegulaFalsiMethod, + ) + dst[2, i] = air_temperature(ts_ρpq) + end +end + + +@testset "Thermodynamics - kernels" begin + FT = Float32 + dev = device(ArrayType) + profiles = PhaseEquilProfiles(param_set, Array) + slice = Colon() + profiles = convert_profile_set(profiles, ArrayType, slice) + + n_profiles = length(profiles.z) + n_vars = length(propertynames(profiles)) + d_dst = ArrayType(Array{FT}(undef, 2, n_profiles)) + fill!(d_dst, 0) + + @unpack e_int, ρ, p, q_tot = profiles + + work_groups = (1,) + ndrange = (n_profiles,) + kernel! = test_thermo_kernel!(dev, work_groups) + event = kernel!(param_set, d_dst, e_int, ρ, p, q_tot, ndrange = ndrange) + wait(dev, event) + + ts_correct = PhaseEquil.(param_set, Array(e_int), Array(ρ), Array(q_tot)) + @test all(Array(d_dst)[1, :] .≈ air_temperature.(ts_correct)) + + ts_correct = + PhaseEquil_ρpq.( + param_set, + Array(ρ), + Array(p), + Array(q_tot), + true, + 100, + RegulaFalsiMethod, + ) + @test all(Array(d_dst)[2, :] .≈ air_temperature.(ts_correct)) + +end