Merge #2273

2273: Move temporary variables in cache in `p.scratch` r=Sbozzolo a=Sbozzolo

To be merged after #2270


Co-authored-by: Gabriele Bozzola <gbozzola@caltech.edu>

src/cache/cache.jl

@@ -93,7 +93,7 @@ function default_cache(
         net_energy_flux_sfc,
         env_thermo_quad = SGSQuadrature(FT),
         precomputed_quantities(Y, atmos)...,
-        temporary_quantities(atmos, spaces.center_space, spaces.face_space)...,
+        scratch = temporary_quantities(Y, atmos),
         hyperdiffusion_cache(Y, atmos, do_dss)...,
     )
     set_precomputed_quantities!(Y, default_cache, FT(0))

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -272,13 +272,13 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
     thermo_params = CAP.thermodynamics_params(params)
     microphys_params = CAP.microphysics_params(params)
 
-    ᶜ∇Φ³ = p.ᶜtemp_CT3
+    ᶜ∇Φ³ = p.scratch.ᶜtemp_CT3
     @. ᶜ∇Φ³ = CT3(ᶜgradᵥ(ᶠinterp(ᶜΦ)))
     @. ᶜ∇Φ³ += CT3(gradₕ(ᶜΦ))
 
-    ρaʲu³ʲ_data = p.temp_data_level
-    u³ʲ_datau³ʲ_data = p.temp_data_level_2
-    ρaʲu³ʲ_datah_tot = ρaʲu³ʲ_dataq_tot = p.temp_data_level_3
+    ρaʲu³ʲ_data = p.scratch.temp_data_level
+    u³ʲ_datau³ʲ_data = p.scratch.temp_data_level_2
+    ρaʲu³ʲ_datah_tot = ρaʲu³ʲ_dataq_tot = p.scratch.temp_data_level_3
 
     z_sfc_halflevel =
         Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
@@ -725,20 +725,20 @@ function set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t)
     )
 
     # TODO: Currently the shear production only includes vertical gradients
-    ᶠu⁰ = p.ᶠtemp_C123
+    ᶠu⁰ = p.scratch.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
-    ᶜstrain_rate = p.ᶜtemp_UVWxUVW
+    ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
     compute_strain_rate_center!(ᶜstrain_rate, ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
-    ᶜprandtl_nvec = p.ᶜtemp_scalar
+    ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
     @. ᶜprandtl_nvec = turbulent_prandtl_number(
         params,
         obukhov_length,
         ᶜlinear_buoygrad,
         ᶜstrain_rate_norm,
     )
-    ᶜtke_exch = p.ᶜtemp_scalar_2
+    ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
     @. ᶜtke_exch = 0
     @. ᶜtke⁰ = Y.c.sgs⁰.ρatke / Y.c.ρ
     # using ᶜu⁰ would be more correct, but this is more consistent with the

src/cache/precomputed_quantities.jl

@@ -302,7 +302,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     n = n_mass_flux_subdomains(turbconv_model)
     thermo_args = (thermo_params, energy_form, moisture_model)
     (; ᶜspecific, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp, ᶜΦ) = p
-    ᶠuₕ³ = p.ᶠtemp_CT3
+    ᶠuₕ³ = p.scratch.ᶠtemp_CT3
 
     @. ᶜspecific = specific_gs(Y.c)
     set_ᶠuₕ³!(ᶠuₕ³, Y)
@@ -364,8 +364,8 @@ values of the first updraft.
 function output_prognostic_sgs_quantities(Y, p, t)
     (; turbconv_model) = p.atmos
     thermo_params = CAP.thermodynamics_params(p.params)
-    (; ᶜp, ᶜρa⁰, ᶜρ⁰, ᶜΦ, ᶜtsʲs) = p
-    ᶠuₕ³ = p.ᶠtemp_CT3
+    (; ᶜρa⁰, ᶜρ⁰, ᶜtsʲs) = p
+    ᶠuₕ³ = p.scratch.ᶠtemp_CT3
     set_ᶠuₕ³!(ᶠuₕ³, Y)
     (ᶠu₃⁺, ᶜu⁺, ᶠu³⁺, ᶜK⁺) = similar.((p.ᶠu₃⁰, p.ᶜu⁰, p.ᶠu³⁰, p.ᶜK⁰))
     set_sgs_ᶠu₃!(u₃⁺, ᶠu₃⁺, Y, turbconv_model)

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -244,20 +244,20 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
     )
 
     # TODO: Currently the shear production only includes vertical gradients
-    ᶠu⁰ = p.ᶠtemp_C123
+    ᶠu⁰ = p.scratch.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
-    ᶜstrain_rate = p.ᶜtemp_UVWxUVW
+    ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
     compute_strain_rate_center!(ᶜstrain_rate, ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
-    ᶜprandtl_nvec = p.ᶜtemp_scalar
+    ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
     @. ᶜprandtl_nvec = turbulent_prandtl_number(
         params,
         obukhov_length,
         ᶜlinear_buoygrad,
         ᶜstrain_rate_norm,
     )
-    ᶜtke_exch = p.ᶜtemp_scalar_2
+    ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
     @. ᶜtke_exch = 0
     for j in 1:n
         @. ᶜtke_exch +=

src/cache/temporary_quantities.jl

@@ -1,35 +1,15 @@
-using LinearAlgebra: ×, norm, dot
-
-import .Parameters as CAP
-using ClimaCore: Operators, Fields, Limiters, Geometry, Spaces
-
-import ClimaComms
-using ClimaCore.Geometry: ⊗
-
-import Thermodynamics as TD
-
+using ClimaCore: Fields
 using ClimaCore.Utilities: half
 
-import ClimaCore.Fields: ColumnField
-
-# Functions on which the model depends:
-# CAP.R_d(params)         # dry specific gas constant
-# CAP.kappa_d(params)     # dry adiabatic exponent
-# CAP.T_triple(params)    # triple point temperature of water
-# CAP.MSLP(params)        # reference pressure
-# CAP.grav(params)        # gravitational acceleration
-# CAP.Omega(params)       # rotation rate (only used if space is spherical)
-# CAP.cv_d(params)        # dry isochoric specific heat capacity
-# The value of cv_d is implied by the values of R_d and kappa_d
-
 # The model also depends on f_plane_coriolis_frequency(params)
 # This is a constant Coriolis frequency that is only used if space is flat
 
 # Fields used to store variables that only need to be used in a single function
 # but cannot be computed on the fly. Unlike the precomputed quantities, these
 # can be modified at any point, so they should never be assumed to be unchanged
 # between function calls.
-function temporary_quantities(atmos, center_space, face_space)
+function temporary_quantities(Y, atmos)
+    center_space, face_space = axes(Y.c), axes(Y.f)
     FT = Spaces.undertype(center_space)
     n = n_mass_flux_subdomains(atmos.turbconv_model)
     return (;

src/prognostic_equations/advection.jl

@@ -84,10 +84,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     ᶜρa⁰ = advect_tke ? (n > 0 ? p.ᶜρa⁰ : Y.c.ρ) : nothing
     ᶜρ⁰ = advect_tke ? (n > 0 ? p.ᶜρ⁰ : Y.c.ρ) : nothing
     ᶜtke⁰ = advect_tke ? p.ᶜtke⁰ : nothing
-    ᶜa_scalar = p.ᶜtemp_scalar
-    ᶜω³ = p.ᶜtemp_CT3
-    ᶠω¹² = p.ᶠtemp_CT12
-    ᶠω¹²ʲs = p.ᶠtemp_CT12ʲs
+    ᶜa_scalar = p.scratch.ᶜtemp_scalar
+    ᶜω³ = p.scratch.ᶜtemp_CT3
+    ᶠω¹² = p.scratch.ᶠtemp_CT12
+    ᶠω¹²ʲs = p.scratch.ᶠtemp_CT12ʲs
 
     if point_type <: Geometry.Abstract3DPoint
         @. ᶜω³ = curlₕ(Y.c.uₕ)

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -24,8 +24,8 @@ function edmfx_sgs_mass_flux_tendency!(
 
     if p.atmos.edmfx_sgs_mass_flux
         # energy
-        ᶠu³_diff_colidx = p.ᶠtemp_CT3[colidx]
-        ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.ᶜtemp_scalar[colidx]
+        ᶠu³_diff_colidx = p.scratch.ᶠtemp_CT3[colidx]
+        ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.scratch.ᶜtemp_scalar[colidx]
         for j in 1:n
             @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
             @. ᶜh_tot_diff_colidx =
@@ -109,8 +109,8 @@ function edmfx_sgs_mass_flux_tendency!(
 
     if p.atmos.edmfx_sgs_mass_flux
         # energy
-        ᶠu³_diff_colidx = p.ᶠtemp_CT3[colidx]
-        ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.ᶜtemp_scalar[colidx]
+        ᶠu³_diff_colidx = p.scratch.ᶠtemp_CT3[colidx]
+        ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.scratch.ᶜtemp_scalar[colidx]
         for j in 1:n
             @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
             @. ᶜh_tot_diff_colidx = ᶜh_totʲs.:($$j)[colidx] - ᶜh_tot[colidx]
@@ -169,7 +169,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     if p.atmos.edmfx_sgs_diffusive_flux
         # energy
-        ᶠρaK_h = p.ᶠtemp_scalar
+        ᶠρaK_h = p.scratch.ᶠtemp_scalar
         @. ᶠρaK_h[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_h[colidx])
 
         ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(
@@ -192,9 +192,9 @@ function edmfx_sgs_diffusive_flux_tendency!(
         end
 
         # momentum
-        ᶠρaK_u = p.ᶠtemp_scalar
+        ᶠρaK_u = p.scratch.ᶠtemp_scalar
         @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx])
-        ᶠstrain_rate = p.ᶠtemp_UVWxUVW
+        ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
         compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu⁰[colidx])
         @. Yₜ.c.uₕ[colidx] -= C12(
             ᶜdivᵥ(-(2 * ᶠρaK_u[colidx] * ᶠstrain_rate[colidx])) / Y.c.ρ[colidx],
@@ -235,7 +235,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     if p.atmos.edmfx_sgs_diffusive_flux
         # energy
-        ᶠρaK_h = p.ᶠtemp_scalar
+        ᶠρaK_h = p.scratch.ᶠtemp_scalar
         @. ᶠρaK_h[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_h[colidx])
 
         ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(
@@ -259,9 +259,9 @@ function edmfx_sgs_diffusive_flux_tendency!(
         end
 
         # momentum
-        ᶠρaK_u = p.ᶠtemp_scalar
+        ᶠρaK_u = p.scratch.ᶠtemp_scalar
         @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx])
-        ᶠstrain_rate = p.ᶠtemp_UVWxUVW
+        ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
         compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu[colidx])
         @. Yₜ.c.uₕ[colidx] -= C12(
             ᶜdivᵥ(-(2 * ᶠρaK_u[colidx] * ᶠstrain_rate[colidx])) / Y.c.ρ[colidx],

src/prognostic_equations/edmfx_tke.jl

@@ -23,7 +23,7 @@ function edmfx_tke_tendency!(
     (; ᶜK_u, ᶜK_h, ρatke_flux) = p
     ᶠgradᵥ = Operators.GradientC2F()
 
-    ᶠρaK_u = p.ᶠtemp_scalar
+    ᶠρaK_u = p.scratch.ᶠtemp_scalar
     if use_prognostic_tke(turbconv_model)
         # turbulent transport (diffusive flux)
         @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx])

src/prognostic_equations/implicit/implicit_solver.jl

@@ -217,7 +217,7 @@ function Wfact!(A, Y, p, dtγ, t)
             p.ᶠgradᵥ_ᶜΦ,
             p.ᶜρ_ref,
             p.ᶜp_ref,
-            p.ᶜtemp_scalar,
+            p.scratch.ᶜtemp_scalar,
             p.params,
             p.atmos,
             (energy_form isa TotalEnergy ? (; p.ᶜh_tot) : (;))...,

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -65,9 +65,9 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
     (; dt) = p.simulation
     n = n_mass_flux_subdomains(turbconv_model)
     ᶜJ = Fields.local_geometry_field(Y.c).J
-    (; ᶜspecific, ᶠu³, ᶜp, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, ᶜtemp_scalar) = p
+    (; ᶜspecific, ᶠu³, ᶜp, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p
 
-    ᶜ1 = ᶜtemp_scalar
+    ᶜ1 = p.scratch.ᶜtemp_scalar
     @. ᶜ1[colidx] = one(Y.c.ρ[colidx])
     vertical_transport!(
         Yₜ.c.ρ[colidx],

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -84,7 +84,7 @@ function vertical_diffusion_boundary_layer_tendency!(
 
     FT = eltype(Y)
     interior_uₕ = Fields.level(Y.c.uₕ, 1)
-    ᶠp = ᶠρK_E = p.ᶠtemp_scalar
+    ᶠp = ᶠρK_E = p.scratch.ᶠtemp_scalar
     @. ᶠp[colidx] = ᶠinterp(ᶜp[colidx])
     ᶜΔz_surface = Fields.Δz_field(interior_uₕ)
     @. ᶠρK_E[colidx] =
@@ -114,8 +114,8 @@ function vertical_diffusion_boundary_layer_tendency!(
         @. Yₜ.c.ρe_tot[colidx] -=
             ᶜdivᵥ_ρe_tot(-(ᶠρK_E[colidx] * ᶠgradᵥ(ᶜh_tot[colidx])))
     end
-    ᶜρχₜ_diffusion = p.ᶜtemp_scalar
-    ρ_flux_χ = p.sfc_temp_C3
+    ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar
+    ρ_flux_χ = p.scratch.sfc_temp_C3
     for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
         χ_name == :e_tot && continue
         if χ_name == :q_tot

src/solver/type_getters.jl

@@ -99,7 +99,6 @@ function get_numerics(parsed_args)
     limiter =
         parsed_args["apply_limiter"] ? Limiters.QuasiMonotoneLimiter : nothing
 
-
     # wrap each upwinding mode in a Val for dispatch
     numerics = AtmosNumerics(;
         energy_upwinding,

src/surface_conditions/surface_conditions.jl

@@ -107,7 +107,8 @@ This functions needs to be called by the coupler whenever either field changes
 to ensure that the simulation is properly updated.
 """
 function set_surface_conditions!(p, surface_conditions, surface_ts)
-    (; sfc_conditions, params, atmos, ᶠtemp_scalar) = p
+    (; sfc_conditions, params, atmos) = p
+    (; ᶠtemp_scalar) = p.scratch
 
     FT = eltype(params)
     FT′ = eltype(parent(surface_conditions))