Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Move temporary variables in cache in p.scratch #2273

Merged
merged 3 commits into from
Oct 25, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/cache/cache.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
16 changes: 8 additions & 8 deletions src/cache/diagnostic_edmf_precomputed_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions src/cache/precomputed_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 4 additions & 4 deletions src/cache/prognostic_edmf_precomputed_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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 +=
Expand Down
26 changes: 3 additions & 23 deletions src/cache/temporary_quantities.jl
Original file line number Diff line number Diff line change
@@ -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 (;
Expand Down
8 changes: 4 additions & 4 deletions src/prognostic_equations/advection.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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ₕ)
Expand Down
20 changes: 10 additions & 10 deletions src/prognostic_equations/edmfx_sgs_flux.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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(
Expand All @@ -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],
Expand Down Expand Up @@ -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(
Expand All @@ -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],
Expand Down
2 changes: 1 addition & 1 deletion src/prognostic_equations/edmfx_tke.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
2 changes: 1 addition & 1 deletion src/prognostic_equations/implicit/implicit_solver.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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) : (;))...,
Expand Down
4 changes: 2 additions & 2 deletions src/prognostic_equations/implicit/implicit_tendency.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down
6 changes: 3 additions & 3 deletions src/prognostic_equations/vertical_diffusion_boundary_layer.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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] =
Expand Down Expand Up @@ -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
Expand Down
1 change: 0 additions & 1 deletion src/solver/type_getters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion src/surface_conditions/surface_conditions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
Loading