Eliminate forcing cache

CliMA · Nov 21, 2024 · a3e8058 · a3e8058
1 parent 6fa62ee
commit a3e8058
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 42 deletions.
diff --git a/src/cache/cache.jl b/src/cache/cache.jl
@@ -21,7 +21,6 @@ struct AtmosCache{
     LSAD,
     EXTFORCING,
     EDMFCOR,
-    FOR,
     NONGW,
     ORGW,
     RAD,
@@ -84,7 +83,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
@@ -192,7 +190,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...)
@@ -221,7 +218,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,

diff --git a/src/parameterized_tendencies/radiation/held_suarez.jl b/src/parameterized_tendencies/radiation/held_suarez.jl
@@ -5,29 +5,16 @@
 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 +32,68 @@ function held_suarez_ΔT_y_T_equator(
     return ΔT_y, T_equator
 end
 
+height_factor(σ, σ_b) = max(0, (σ - σ_b) / (1 - σ_b))
+
+struct HeldSuarezForcingParams{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,
+    ts_surf,
+    ρ,
+    p,
+    lat,
+    z_surface,
+    s::HeldSuarezForcingParams,
+)
+    σ = 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 - Δ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,
+    z_surface,
+    p,
+    ts_surf,
+    s::HeldSuarezForcingParams,
+)
+    p / (
+        s.MSLP * exp(
+            -s.grav * z_surface / s.R_d /
+            TD.air_temperature(thermo_params, ts_surf),
+        )
+    )
+end
+
 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 +108,42 @@ 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}(
+        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(
+                compute_σ(thermo_params, z_surface, ᶜp, ts_surf, hs_params),
+                σ_b,
             )
-        )
-
-    @. 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