Merge pull request #3449 from CliMA/ck/elim_cache3

Eliminate forcing cache

Author: charleskawczynski

src/cache/cache.jl

@@ -18,7 +18,6 @@ struct AtmosCache{
     LSAD,
     EXTFORCING,
     EDMFCOR,
-    FOR,
     NONGW,
     ORGW,
     RAD,
@@ -78,7 +77,6 @@ struct AtmosCache{
     large_scale_advection::LSAD
     external_forcing::EXTFORCING
     edmf_coriolis::EDMFCOR
-    forcing::FOR
     non_orographic_gravity_wave::NONGW
     orographic_gravity_wave::ORGW
     radiation::RAD
@@ -183,7 +181,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
     large_scale_advection = large_scale_advection_cache(Y, atmos)
     external_forcing = external_forcing_cache(Y, atmos, params)
     edmf_coriolis = edmf_coriolis_cache(Y, atmos)
-    forcing = forcing_cache(Y, atmos)
     non_orographic_gravity_wave = non_orographic_gravity_wave_cache(Y, atmos)
     orographic_gravity_wave = orographic_gravity_wave_cache(Y, atmos)
     radiation = radiation_model_cache(Y, atmos, radiation_args...)
@@ -209,7 +206,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
         large_scale_advection,
         external_forcing,
         edmf_coriolis,
-        forcing,
         non_orographic_gravity_wave,
         orographic_gravity_wave,
         radiation,

src/parameterized_tendencies/radiation/held_suarez.jl

@@ -2,32 +2,21 @@
 ##### Held-Suarez
 #####
 
+import Thermodynamics as TD
+import Thermodynamics.Parameters as TDP
 import ClimaCore.Spaces as Spaces
 import ClimaCore.Fields as Fields
 
-forcing_cache(Y, atmos::AtmosModel) = forcing_cache(Y, atmos.forcing_type)
-
 #####
 ##### No forcing
 #####
 
-forcing_cache(Y, forcing_type::Nothing) = (;)
 forcing_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 
 #####
 ##### Held-Suarez forcing
 #####
 
-function forcing_cache(Y, forcing_type::HeldSuarezForcing)
-    FT = Spaces.undertype(axes(Y.c))
-    return (;
-        ᶜσ = similar(Y.c, FT),
-        ᶜheight_factor = similar(Y.c, FT),
-        ᶜΔρT = similar(Y.c, FT),
-        ᶜφ = deg2rad.(Fields.coordinate_field(Y.c).lat),
-    )
-end
-
 function held_suarez_ΔT_y_T_equator(params, moisture_model::DryModel)
     FT = eltype(params)
     ΔT_y = FT(CAP.ΔT_y_dry(params))
@@ -45,10 +34,76 @@ function held_suarez_ΔT_y_T_equator(
     return ΔT_y, T_equator
 end
 
+struct HeldSuarezForcingParams{FT}
+    ΔT_y::FT
+    day::FT
+    σ_b::FT
+    R_d::FT
+    T_min::FT
+    T_equator::FT
+    Δθ_z::FT
+    p_ref_theta::FT
+    κ_d::FT
+    grav::FT
+    MSLP::FT
+end
+Base.Broadcast.broadcastable(x::HeldSuarezForcingParams) = tuple(x)
+
+function compute_ΔρT(
+    thermo_params::TDP.ThermodynamicsParameters,
+    ts_surf::TD.ThermodynamicState,
+    ρ::FT,
+    p::FT,
+    lat::FT,
+    z_surface::FT,
+    s::HeldSuarezForcingParams,
+) where {FT}
+    σ = compute_σ(thermo_params, z_surface, p, ts_surf, s)
+    k_a = 1 / (40 * s.day)
+    k_s = 1 / (4 * s.day)
+
+    φ = deg2rad(lat)
+    return (k_a + (k_s - k_a) * height_factor(σ, s.σ_b) * abs2(abs2(cos(φ)))) *
+           ρ *
+           ( # ᶜT - ᶜT_equil
+               p / (ρ * s.R_d) - max(
+                   s.T_min,
+                   (
+                       s.T_equator - s.ΔT_y * abs2(sin(φ)) -
+                       s.Δθ_z * log(p / s.p_ref_theta) * abs2(cos(φ))
+                   ) * fast_pow(p / s.p_ref_theta, s.κ_d),
+               )
+           )
+end
+
+function compute_σ(
+    thermo_params::TDP.ThermodynamicsParameters,
+    z_surface::FT,
+    p::FT,
+    ts_surf::TD.ThermodynamicState,
+    s::HeldSuarezForcingParams,
+) where {FT}
+    p / (
+        s.MSLP * exp(
+            -s.grav * z_surface / s.R_d /
+            TD.air_temperature(thermo_params, ts_surf),
+        )
+    )
+end
+
+height_factor(σ::FT, σ_b::FT) where {FT} = max(0, (σ - σ_b) / (1 - σ_b))
+height_factor(
+    thermo_params::TDP.ThermodynamicsParameters,
+    z_surface::FT,
+    p::FT,
+    ts_surf::TD.ThermodynamicState,
+    s::HeldSuarezForcingParams,
+) where {FT} =
+    height_factor(compute_σ(thermo_params, z_surface, p, ts_surf, s), s.σ_b)
+
 function forcing_tendency!(Yₜ, Y, p, t, ::HeldSuarezForcing)
     (; params) = p
     (; ᶜp, sfc_conditions) = p.precomputed
-    (; ᶜσ, ᶜheight_factor, ᶜΔρT, ᶜφ) = p.forcing
 
     # TODO: Don't need to enforce FT here, it should be done at param creation.
     FT = Spaces.undertype(axes(Y.c))
@@ -63,37 +118,46 @@ function forcing_tendency!(Yₜ, Y, p, t, ::HeldSuarezForcing)
     T_min = FT(CAP.T_min_hs(params))
     thermo_params = CAP.thermodynamics_params(params)
     σ_b = CAP.σ_b(params)
-    k_a = 1 / (40 * day)
-    k_s = 1 / (4 * day)
     k_f = 1 / day
 
     z_surface = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)
 
     ΔT_y, T_equator = held_suarez_ΔT_y_T_equator(params, p.atmos.moisture_model)
 
-    @. ᶜσ =
-        ᶜp / (
-            MSLP * exp(
-                -grav * z_surface / R_d /
-                TD.air_temperature(thermo_params, sfc_conditions.ts),
-            )
-        )
+    hs_params = HeldSuarezForcingParams{FT}(
+        ΔT_y,
+        day,
+        σ_b,
+        R_d,
+        T_min,
+        T_equator,
+        Δθ_z,
+        p_ref_theta,
+        κ_d,
+        grav,
+        MSLP,
+    )
 
-    @. ᶜheight_factor = max(0, (ᶜσ - σ_b) / (1 - σ_b))
-    @. ᶜΔρT =
-        (k_a + (k_s - k_a) * ᶜheight_factor * abs2(abs2(cos(ᶜφ)))) *
-        Y.c.ρ *
-        ( # ᶜT - ᶜT_equil
-            ᶜp / (Y.c.ρ * R_d) - max(
-                T_min,
-                (
-                    T_equator - ΔT_y * abs2(sin(ᶜφ)) -
-                    Δθ_z * log(ᶜp / p_ref_theta) * abs2(cos(ᶜφ))
-                ) * fast_pow(ᶜp / p_ref_theta, κ_d),
+    lat = Fields.coordinate_field(Y.c).lat
+    @. Yₜ.c.uₕ -=
+        (
+            k_f * height_factor(
+                thermo_params,
+                z_surface,
+                ᶜp,
+                sfc_conditions.ts,
+                hs_params,
             )
-        )
-
-    @. Yₜ.c.uₕ -= (k_f * ᶜheight_factor) * Y.c.uₕ
-    @. Yₜ.c.ρe_tot -= ᶜΔρT * cv_d
+        ) * Y.c.uₕ
+    @. Yₜ.c.ρe_tot -=
+        compute_ΔρT(
+            thermo_params,
+            sfc_conditions.ts,
+            Y.c.ρ,
+            ᶜp,
+            lat,
+            z_surface,
+            hs_params,
+        ) * cv_d
     return nothing
 end

test/coupler_compatibility.jl

@@ -84,7 +84,6 @@ const T2 = 290
         p.large_scale_advection,
         p.external_forcing,
         p.edmf_coriolis,
-        p.forcing,
         p.non_orographic_gravity_wave,
         p.orographic_gravity_wave,
         p.radiation,