Merge pull request #2463 from CliMA/ts/thermo-update

Update Thermodynamics API usage

Author: tapios

benchmarks/scripts/thermo_bench.jl

@@ -38,10 +38,13 @@ using JET
 import ClimaCore: Spaces, Fields
 import ClimaCore.Domains: Geometry
 
+@inline function make_thermo_state(thermo_params, ρ, e_int, q_tot)
+    sa = TD.saturation_adjustment(thermo_params, TD.ρe(), ρ, e_int, q_tot; maxiter = 3)
+    p = TD.air_pressure(thermo_params, sa.T, ρ, q_tot, sa.q_liq, sa.q_ice)
+    return (; ρ = ρ, p = p, T = sa.T, e_int = e_int, q_tot = q_tot)
+end
 @inline ts_gs(thermo_params, e_tot, q_tot, K, Φ, ρ) =
-    thermo_state(thermo_params, e_tot - K - Φ, q_tot, ρ)
-@inline thermo_state(thermo_params, ρ, e_int, q_tot) =
-    TD.PhaseEquil_ρeq(thermo_params,ρ,e_int,q_tot, 3, eltype(thermo_params)(0.003))
+    make_thermo_state(thermo_params, ρ, e_tot - K - Φ, q_tot)
 
 import Thermodynamics as TD
 
@@ -115,7 +118,7 @@ using ClimaCore
 import ClimaCore: Spaces, Fields
 import ClimaCore.Domains: Geometry
 
-ENV["CLIMACOMMS_DEVICE"] = "CUDA";
+ENV["CLIMACOMMS_DEVICE"] = get(ENV, "CLIMACOMMS_DEVICE", "CPU");
 ClimaComms.@import_required_backends
 using BenchmarkTools
 @isdefined(TU) || include(
@@ -149,7 +152,7 @@ using Test
         q_tot = FT(0),
         T = FT(0),
     )
-    x = fill((; ts = zero(TD.PhaseEquil{FT}), nt_core...), cspace)
+    x = fill((; ts = nt_ts, nt_core...), cspace)
     xv = fill((; ts = nt_ts, nt_core...), cspace)
     (_, Nij, _, Nv, Nh) = size(Fields.field_values(x.ts))
     us = TB.UniversalSizesStatic(Nv, Nij, Nh)

test/Operators/finitedifference/unit_columnwise.jl

@@ -139,11 +139,6 @@ function implicit_tendency_bc!(Yₜ, Y, p, t)
     return nothing
 end
 
-function thermo_state(thermo_params, ᶜρ, ᶜρe_tot, ᶜK, grav, ᶜz)
-    return @. lazy(
-        TD.PhaseDry_ρe(thermo_params, ᶜρ, ᶜρe_tot / ᶜρ - ᶜK - Φ(grav, ᶜz)),
-    )
-end
 
 function ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t)
     (; rayleigh_sponge, dt, zmax, thermo_params) = p
@@ -158,10 +153,11 @@ function ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t)
     ᶠu₃ = ᶠY.u₃
 
     ᶜK = compute_kinetic(ᶜuₕ, ᶠu₃)
-    ᶜts = thermo_state(thermo_params, ᶜρ, ᶜρe_tot, ᶜK, grav, ᶜz)
-    ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts))
+    ᶜe_int = @. lazy(ᶜρe_tot / ᶜρ - ᶜK - Φ(grav, ᶜz))
+    ᶜT = @. lazy(TD.air_temperature(thermo_params, ᶜe_int))
+    ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜT, ᶜρ))
     ᶜh_tot =
-        @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜρe_tot / ᶜρ))
+        @. lazy(TD.total_enthalpy(thermo_params, ᶜρe_tot / ᶜρ, ᶜT))
     # Central advection of active tracers (e_tot and q_tot)
     ᶠuₕ³ = @. lazy(ᶠwinterp(ᶜρ * ᶜJ, CT3(ᶜuₕ)))
     ᶠu³ = @. lazy(ᶠuₕ³ + CT3(ᶠu₃))
@@ -183,8 +179,9 @@ function ᶠimplicit_tendency_bc(ᶜY, ᶠY, p, t)
     ᶜuₕ = ᶜY.uₕ
     ᶠu₃ = ᶠY.u₃
     ᶜK = compute_kinetic(ᶜuₕ, ᶠu₃)
-    ᶜts = thermo_state(thermo_params, ᶜρ, ᶜρe_tot, ᶜK, grav, ᶜz)
-    ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts))
+    ᶜe_int = @. lazy(ᶜρe_tot / ᶜρ - ᶜK - Φ(grav, ᶜz))
+    ᶜT = @. lazy(TD.air_temperature(thermo_params, ᶜe_int))
+    ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜT, ᶜρ))
     bc1 = @. lazy(-(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) + ᶠgradᵥ(Φ(grav, ᶜz))))
     bc2 = @. lazy(-β_rayleigh_w(rayleigh_sponge, ᶠz, zmax) * ᶠu₃)
     return @. lazy(ᶠtendencies(bc1 + bc2))
@@ -199,7 +196,7 @@ end
 
 function set_precomputed_quantities!(Y, p, t)
     (; thermo_params) = p
-    (; ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp) = p.precomputed
+    (; ᶜu, ᶠu³, ᶠu, ᶜK, ᶜT, ᶜp) = p.precomputed
 
     ᶜρ = Y.c.ρ
     ᶜuₕ = Y.c.uₕ
@@ -211,16 +208,13 @@ function set_precomputed_quantities!(Y, p, t)
     ᶠu³ .= compute_ᶠuₕ³(ᶜuₕ, ᶜρ) .+ CT3.(ᶠu₃)
     ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
 
-    @. ᶜts = TD.PhaseDry_ρe(
-        thermo_params,
-        Y.c.ρ,
-        Y.c.ρe_tot / Y.c.ρ - ᶜK - Φ(grav, ᶜz),
-    )
-    @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
+    ᶜe_int = @. Y.c.ρe_tot / Y.c.ρ - ᶜK - Φ(grav, ᶜz)
+    @. ᶜT = TD.air_temperature(thermo_params, ᶜe_int)
+    @. ᶜp = TD.air_pressure(thermo_params, ᶜT, Y.c.ρ)
 
     (; ᶜh_tot) = p.precomputed
     @. ᶜh_tot =
-        TD.total_specific_enthalpy(thermo_params, ᶜts, Y.c.ρe_tot / Y.c.ρ)
+        TD.total_enthalpy(thermo_params, Y.c.ρe_tot / Y.c.ρ, ᶜT)
     return nothing
 end
 
@@ -300,7 +294,7 @@ precomputed = (;
     ᶠu³ = Fields.Field(CT3{FT}, ᶠspace),
     ᶜp = Fields.Field(FT, ᶜspace),
     ᶜK = Fields.Field(FT, ᶜspace),
-    ᶜts = Fields.Field(TD.PhaseDry{FT}, ᶜspace),
+    ᶜT = Fields.Field(FT, ᶜspace),
     ᶠu = Fields.Field(C123{FT}, ᶠspace),
     ᶜu = Fields.Field(C123{FT}, ᶜspace),
 )