From 555b9a22de8345e0890b2984be1c5e1efe8cda2d Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Wed, 25 Oct 2023 10:54:25 -0700 Subject: [PATCH 1/3] Cache: move temporary quantities to scratch --- src/cache/cache.jl | 2 +- .../diagnostic_edmf_precomputed_quantities.jl | 16 +++++++-------- src/cache/precomputed_quantities.jl | 6 +++--- .../prognostic_edmf_precomputed_quantities.jl | 8 ++++---- src/cache/temporary_quantities.jl | 3 ++- src/prognostic_equations/advection.jl | 8 ++++---- src/prognostic_equations/edmfx_sgs_flux.jl | 20 +++++++++---------- src/prognostic_equations/edmfx_tke.jl | 2 +- .../implicit/implicit_solver.jl | 2 +- .../implicit/implicit_tendency.jl | 4 ++-- .../vertical_diffusion_boundary_layer.jl | 6 +++--- src/solver/type_getters.jl | 1 - src/surface_conditions/surface_conditions.jl | 3 ++- 13 files changed, 41 insertions(+), 40 deletions(-) diff --git a/src/cache/cache.jl b/src/cache/cache.jl index 6439132353..02ea1dbfea 100644 --- a/src/cache/cache.jl +++ b/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)) diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 2137b331f0..85db783f71 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/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 diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 18d88e6b30..56901bdc14 100644 --- a/src/cache/precomputed_quantities.jl +++ b/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) diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 9079674423..f472390d1d 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/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 += diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl index e43a05acda..bc4ad2b630 100644 --- a/src/cache/temporary_quantities.jl +++ b/src/cache/temporary_quantities.jl @@ -29,7 +29,8 @@ import ClimaCore.Fields: ColumnField # 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 (; diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 992d9db88c..c3a500edcc 100644 --- a/src/prognostic_equations/advection.jl +++ b/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ₕ) diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index c2e861fe9e..47a3b957e1 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/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], diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index 07841b25db..3676ae0e9e 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/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]) diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index 6b03fea548..49e28aa642 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/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) : (;))..., diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 6c73ef4f08..e62070b1bc 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/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], diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index c00963d1ee..dc68cc3280 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/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 diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 0387475668..2e843a057f 100644 --- a/src/solver/type_getters.jl +++ b/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, diff --git a/src/surface_conditions/surface_conditions.jl b/src/surface_conditions/surface_conditions.jl index 5bbf21284d..1cd137b89e 100644 --- a/src/surface_conditions/surface_conditions.jl +++ b/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)) From 0fa38c312a3f3e07f7a9a8576d4458ae7a5f2fb4 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Wed, 25 Oct 2023 10:57:20 -0700 Subject: [PATCH 2/3] Remove unnecessry comments --- src/cache/temporary_quantities.jl | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl index bc4ad2b630..48654a8498 100644 --- a/src/cache/temporary_quantities.jl +++ b/src/cache/temporary_quantities.jl @@ -12,16 +12,6 @@ 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 From b85911657690ebaa9887389dd9cb65f1cc6ca195 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Wed, 25 Oct 2023 10:58:30 -0700 Subject: [PATCH 3/3] Remove unncessary imports --- src/cache/temporary_quantities.jl | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl index 48654a8498..77861ba8e2 100644 --- a/src/cache/temporary_quantities.jl +++ b/src/cache/temporary_quantities.jl @@ -1,17 +1,6 @@ -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 - # The model also depends on f_plane_coriolis_frequency(params) # This is a constant Coriolis frequency that is only used if space is flat