From f50a1d325039734e1a211a2e6cbd5b87f1ffdc3e Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Thu, 12 Oct 2023 10:59:53 -0700 Subject: [PATCH] Cache: move temporary quantities to scratch --- src/cache/cache.jl | 6 +++++- .../diagnostic_edmf_precomputed_quantities.jl | 16 +++++++-------- src/cache/precomputed_quantities.jl | 6 +++--- .../prognostic_edmf_precomputed_quantities.jl | 8 ++++---- 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 ++- 12 files changed, 43 insertions(+), 39 deletions(-) diff --git a/src/cache/cache.jl b/src/cache/cache.jl index 6439132353f..7cf30b04cb1 100644 --- a/src/cache/cache.jl +++ b/src/cache/cache.jl @@ -93,7 +93,11 @@ 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( + atmos, + spaces.center_space, + spaces.face_space, + ), 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 80236a82b43..59b07784e82 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -155,13 +155,13 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t) @. ᶜtke⁰ = Y.c.sgs⁰.ρatke / Y.c.ρ - ᶜ∇Φ³ = 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 ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1)) uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1)) @@ -678,20 +678,20 @@ function set_diagnostic_edmf_precomputed_quantities!(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 # using ᶜu⁰ would be more correct, but this is more consistent with the # TKE equation, where using ᶜu⁰ results in allocation diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index bdaf9f47c89..bef3daba645 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) @@ -359,8 +359,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 a0a62358d75..eaaea50fc09 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -207,20 +207,20 @@ function set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, 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/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 992d9db88c2..c3a500edcc4 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 c2e861fe9e5..47a3b957e17 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 07841b25db1..3676ae0e9e0 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 6b03fea5481..49e28aa6421 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 6c73ef4f085..e62070b1bc0 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 c00963d1ee7..dc68cc32800 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 03874756680..2e843a057fa 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 5bbf21284d9..1cd137b89e0 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))