Skip to content

Commit 41acb83

Browse files
costachriscostachris
committed
Remove k_u, k_v, and mixing length from precomputed quantities
1 parent fc1bacd commit 41acb83

14 files changed

+335
-193
lines changed

‎src/cache/cloud_fraction.jl

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,9 @@ end
2222
Compute the grid scale cloud fraction based on sub-grid scale properties
2323
"""
2424
NVTX.@annotate function set_cloud_fraction!(Y, p, ::DryModel, _)
25-
(; ᶜmixing_length) = p.precomputed
2625
(; turbconv_model) = p.atmos
2726
FT = eltype(p.params)
28-
if isnothing(turbconv_model)
29-
compute_gm_mixing_length!(ᶜmixing_length, Y, p)
30-
end
27+
3128
p.precomputed.cloud_diagnostics_tuple .=
3229
((; cf = FT(0), q_liq = FT(0), q_ice = FT(0)),)
3330
end
@@ -39,8 +36,9 @@ NVTX.@annotate function set_cloud_fraction!(
3936
)
4037
(; params) = p
4138
(; turbconv_model) = p.atmos
42-
(; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
39+
(; ᶜts, cloud_diagnostics_tuple) = p.precomputed
4340
thermo_params = CAP.thermodynamics_params(params)
41+
4442
if isnothing(turbconv_model)
4543
if p.atmos.call_cloud_diagnostics_per_stage isa
4644
CallCloudDiagnosticsPerStage
@@ -53,7 +51,6 @@ NVTX.@annotate function set_cloud_fraction!(
5351
@. ᶜgradᵥ_θ_liq_ice =
5452
ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
5553
end
56-
compute_gm_mixing_length!(ᶜmixing_length, Y, p)
5754
end
5855
if moist_model isa EquilMoistModel
5956
@. cloud_diagnostics_tuple = make_named_tuple(
@@ -85,7 +82,7 @@ NVTX.@annotate function set_cloud_fraction!(
8582

8683
FT = eltype(params)
8784
thermo_params = CAP.thermodynamics_params(params)
88-
(; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
85+
(; ᶜts, cloud_diagnostics_tuple) = p.precomputed
8986
(; turbconv_model) = p.atmos
9087

9188
if isnothing(turbconv_model)
@@ -100,9 +97,10 @@ NVTX.@annotate function set_cloud_fraction!(
10097
@. ᶜgradᵥ_θ_liq_ice =
10198
ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
10299
end
103-
compute_gm_mixing_length!(ᶜmixing_length, Y, p)
104100
end
105101

102+
ᶜmixing_length = compute_gm_mixing_length(Y, p)
103+
106104
diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
107105
@. cloud_diagnostics_tuple = quad_loop(
108106
SG_quad,
@@ -139,13 +137,17 @@ NVTX.@annotate function set_cloud_fraction!(
139137

140138
FT = eltype(params)
141139
thermo_params = CAP.thermodynamics_params(params)
142-
(; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
140+
(; ᶜts, cloud_diagnostics_tuple) = p.precomputed
143141
(; turbconv_model) = p.atmos
144142

145143
# TODO - we should make this default when using diagnostic edmf
146144
# environment
147145
diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
148146

147+
ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
148+
ᶜmixing_length_lazy = mixing_length(Y, p)
149+
ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
150+
149151
@. cloud_diagnostics_tuple = quad_loop(
150152
SG_quad,
151153
ᶜts,
@@ -192,14 +194,18 @@ NVTX.@annotate function set_cloud_fraction!(
192194

193195
FT = eltype(params)
194196
thermo_params = CAP.thermodynamics_params(params)
195-
(; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
197+
(; ᶜts⁰, cloud_diagnostics_tuple) = p.precomputed
196198
(; ᶜρʲs, ᶜtsʲs, ᶜρa⁰, ᶜρ⁰) = p.precomputed
197199
(; turbconv_model) = p.atmos
198200

199201
# TODO - we should make this default when using diagnostic edmf
200202
# environment
201203
diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
202204

205+
ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
206+
ᶜmixing_length_lazy = mixing_length(Y, p)
207+
ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
208+
203209
@. cloud_diagnostics_tuple = quad_loop(
204210
SG_quad,
205211
ᶜts⁰,

‎src/cache/diagnostic_edmf_precomputed_quantities.jl

Lines changed: 2 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -963,13 +963,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
963963
(; ᶜp, ᶜu, ᶜts) = p.precomputed
964964
(; ustar, obukhov_length) = p.precomputed.sfc_conditions
965965
(; ᶜtke⁰) = p.precomputed
966-
(;
967-
ᶜlinear_buoygrad,
968-
ᶜstrain_rate_norm,
969-
ᶜmixing_length_tuple,
970-
ᶜmixing_length,
971-
) = p.precomputed
972-
(; ᶜK_h, ᶜK_u, ρatke_flux) = p.precomputed
966+
(; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
967+
(; ρatke_flux) = p.precomputed
973968
turbconv_params = CAP.turbconv_params(params)
974969
thermo_params = CAP.thermodynamics_params(params)
975970
ᶜlg = Fields.local_geometry_field(Y.c)
@@ -1001,42 +996,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
1001996
ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
1002997
@. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
1003998

1004-
ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
1005-
@. ᶜprandtl_nvec =
1006-
turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
1007-
1008-
ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
1009-
@. ᶜtke_exch = 0
1010-
# using ᶜu⁰ would be more correct, but this is more consistent with the
1011-
# TKE equation, where using ᶜu⁰ results in allocation
1012-
for j in 1:n
1013-
@. ᶜtke_exch +=
1014-
ᶜρaʲs.:($$j) * ᶜdetrʲs.:($$j) / Y.c.ρ *
1015-
(1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰)
1016-
end
1017-
1018-
sfc_tke = Fields.level(ᶜtke⁰, 1)
1019-
z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, half)
1020-
@. ᶜmixing_length_tuple = mixing_length(
1021-
params,
1022-
ustar,
1023-
ᶜz,
1024-
z_sfc,
1025-
ᶜdz,
1026-
max(sfc_tke, 0),
1027-
ᶜlinear_buoygrad,
1028-
max(ᶜtke⁰, 0),
1029-
obukhov_length,
1030-
ᶜstrain_rate_norm,
1031-
ᶜprandtl_nvec,
1032-
ᶜtke_exch,
1033-
p.atmos.edmfx_model.scale_blending_method,
1034-
)
1035-
@. ᶜmixing_length = ᶜmixing_length_tuple.master
1036-
1037-
@. ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
1038-
@. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
1039-
1040999
ρatke_flux_values = Fields.field_values(ρatke_flux)
10411000
ρ_int_values = Fields.field_values(Fields.level(Y.c.ρ, 1))
10421001
u_int_values = Fields.field_values(Fields.level(ᶜu, 1))

‎src/cache/precomputed_quantities.jl

Lines changed: 1 addition & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,6 @@ function precomputed_quantities(Y, atmos)
113113
gs_quantities = (;
114114
ᶜwₜqₜ = similar(Y.c, Geometry.WVector{FT}),
115115
ᶜwₕhₜ = similar(Y.c, Geometry.WVector{FT}),
116-
ᶜmixing_length = similar(Y.c, FT),
117116
ᶜlinear_buoygrad = similar(Y.c, FT),
118117
ᶜstrain_rate_norm = similar(Y.c, FT),
119118
sfc_conditions = similar(Spaces.level(Y.f, half), SCT),
@@ -172,9 +171,6 @@ function precomputed_quantities(Y, atmos)
172171
advective_sgs_quantities =
173172
atmos.turbconv_model isa PrognosticEDMFX ?
174173
(;
175-
ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
176-
ᶜK_u = similar(Y.c, FT),
177-
ᶜK_h = similar(Y.c, FT),
178174
ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
179175
bdmr_l = similar(Y.c, BidiagonalMatrixRow{FT}),
180176
bdmr_r = similar(Y.c, BidiagonalMatrixRow{FT}),
@@ -192,11 +188,8 @@ function precomputed_quantities(Y, atmos)
192188
edonly_quantities =
193189
atmos.turbconv_model isa EDOnlyEDMFX ?
194190
(;
195-
ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
196191
ᶜtke⁰ = similar(Y.c, FT),
197-
ᶜK_u = similar(Y.c, FT),
198192
ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
199-
ᶜK_h = similar(Y.c, FT),
200193
) : (;)
201194

202195
sgs_quantities = (;
@@ -224,20 +217,10 @@ function precomputed_quantities(Y, atmos)
224217
ᶠu³⁰ = similar(Y.f, CT3{FT}),
225218
ᶜu⁰ = similar(Y.c, C123{FT}),
226219
ᶜK⁰ = similar(Y.c, FT),
227-
ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
228-
ᶜK_u = similar(Y.c, FT),
229-
ᶜK_h = similar(Y.c, FT),
230220
ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
231221
precipitation_sgs_quantities...,
232222
) : (;)
233-
vert_diff_quantities =
234-
if atmos.vert_diff isa
235-
Union{VerticalDiffusion, DecayWithHeightDiffusion}
236-
ᶜK_h = similar(Y.c, FT)
237-
(; ᶜK_u = ᶜK_h, ᶜK_h) # ᶜK_u aliases ᶜK_h because they are always equal.
238-
else
239-
(;)
240-
end
223+
241224
smagorinsky_lilly_quantities =
242225
if atmos.smagorinsky_lilly isa SmagorinskyLilly
243226
uvw_vec = UVW(FT(0), FT(0), FT(0))
@@ -258,7 +241,6 @@ function precomputed_quantities(Y, atmos)
258241
advective_sgs_quantities...,
259242
edonly_quantities...,
260243
diagnostic_sgs_quantities...,
261-
vert_diff_quantities...,
262244
sedimentation_quantities...,
263245
precipitation_quantities...,
264246
surface_precip_fluxes...,
@@ -534,12 +516,6 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t)
534516
ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
535517
end
536518

537-
# TODO: It is too slow to calculate mixing length at every timestep
538-
# if isnothing(turbconv_model)
539-
# (; ᶜmixing_length) = p.precomputed
540-
# compute_gm_mixing_length!(ᶜmixing_length, Y, p)
541-
# end
542-
543519
if turbconv_model isa PrognosticEDMFX
544520
set_prognostic_edmf_precomputed_quantities_bottom_bc!(Y, p, t)
545521
set_prognostic_edmf_precomputed_quantities_explicit_closures!(Y, p, t)
@@ -576,14 +552,6 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t)
576552
set_precipitation_cache!(Y, p, p.atmos.precip_model, p.atmos.turbconv_model)
577553
set_precipitation_surface_fluxes!(Y, p, p.atmos.precip_model)
578554

579-
if vert_diff isa DecayWithHeightDiffusion
580-
(; ᶜK_h) = p.precomputed
581-
@. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.ρ, vert_diff)
582-
elseif vert_diff isa VerticalDiffusion
583-
(; ᶜK_h) = p.precomputed
584-
@. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vert_diff)
585-
end
586-
587555
# TODO
588556
if call_cloud_diagnostics_per_stage isa CallCloudDiagnosticsPerStage
589557
set_cloud_fraction!(Y, p, moisture_model, cloud_model)

‎src/cache/prognostic_edmf_precomputed_quantities.jl

Lines changed: 1 addition & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -368,15 +368,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
368368
n = n_mass_flux_subdomains(turbconv_model)
369369

370370
(; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
371-
(;
372-
ᶜmixing_length_tuple,
373-
ᶜmixing_length,
374-
ᶜlinear_buoygrad,
375-
ᶜstrain_rate_norm,
376-
ᶜK_u,
377-
ᶜK_h,
378-
ρatke_flux,
379-
) = p.precomputed
371+
(; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ρatke_flux) = p.precomputed
380372
(;
381373
ᶜuʲs,
382374
ᶜtsʲs,
@@ -501,41 +493,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
501493
ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
502494
@. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
503495

504-
ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
505-
@. ᶜprandtl_nvec =
506-
turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
507-
508-
ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
509-
@. ᶜtke_exch = 0
510-
for j in 1:n
511-
ᶠu³ʲ = ᶠu³ʲs.:($j)
512-
@. ᶜtke_exch +=
513-
Y.c.sgsʲs.:($$j).ρa * ᶜdetrʲs.:($$j) / ᶜρa⁰ *
514-
(1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰)
515-
end
516-
517-
sfc_tke = Fields.level(ᶜtke⁰, 1)
518-
@. ᶜmixing_length_tuple = mixing_length(
519-
p.params,
520-
ustar,
521-
ᶜz,
522-
z_sfc,
523-
ᶜdz,
524-
max(sfc_tke, eps(FT)),
525-
ᶜlinear_buoygrad,
526-
max(ᶜtke⁰, 0),
527-
obukhov_length,
528-
ᶜstrain_rate_norm,
529-
ᶜprandtl_nvec,
530-
ᶜtke_exch,
531-
p.atmos.edmfx_model.scale_blending_method,
532-
)
533-
534-
@. ᶜmixing_length = ᶜmixing_length_tuple.master
535-
536-
@. ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
537-
@. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
538-
539496
ρatke_flux_values = Fields.field_values(ρatke_flux)
540497
ρa_sfc_values = Fields.field_values(Fields.level(ᶜρa⁰, 1)) # TODO: replace by surface value
541498
ustar_values = Fields.field_values(ustar)

‎src/cache/temporary_quantities.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ function temporary_quantities(Y, atmos)
1818
ᶜtemp_scalar_3 = Fields.Field(FT, center_space),
1919
ᶜtemp_scalar_4 = Fields.Field(FT, center_space),
2020
ᶜtemp_scalar_5 = Fields.Field(FT, center_space),
21+
ᶜtemp_scalar_6 = Fields.Field(FT, center_space),
22+
ᶜtemp_scalar_7 = Fields.Field(FT, center_space),
2123
ᶠtemp_field_level = Fields.level(Fields.Field(FT, face_space), half),
2224
temp_field_level = Fields.level(Fields.Field(FT, center_space), 1),
2325
temp_field_level_2 = Fields.level(Fields.Field(FT, center_space), 1),

‎src/diagnostics/Diagnostics.jl

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,15 @@ import ..PrognosticSurfaceTemperature
5454

5555
# functions used to calculate diagnostics
5656
import ..draft_area
57-
import ..compute_gm_mixing_length!
57+
import ..compute_gm_mixing_length
5858
import ..horizontal_integral_at_boundary
59+
import ..mixing_length
60+
import ..eddy_diffusivity
61+
import ..eddy_viscosity
62+
import ..turbulent_prandtl_number
63+
import ..smagorinsky_lilly_length
64+
import ..compute_eddy_diffusivity_coefficient
65+
5966

6067
# We need the abbreviations for symbols like curl, grad, and so on
6168
include(joinpath("..", "utils", "abbreviations.jl"))

‎src/diagnostics/core_diagnostics.jl

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -273,10 +273,39 @@ add_diagnostic_variable!(
273273
or from mixing length closure with EDMF SGS model.
274274
""",
275275
compute! = (out, state, cache, time) -> begin
276+
turbconv_model = cache.atmos.turbconv_model
277+
# TODO: consolidate remaining mixing length types
278+
# (smagorinsky_lilly, dz) into a single mixing length function
279+
if isa(turbconv_model, PrognosticEDMFX) ||
280+
isa(turbconv_model, DiagnosticEDMFX) ||
281+
isa(turbconv_model, EDOnlyEDMFX)
282+
ᶜmixing_length = mixing_length(state, cache)
283+
else
284+
(; params) = cache
285+
(; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = cache.precomputed
286+
ᶜdz = Fields.Δz_field(axes(state.c))
287+
ᶜprandtl_nvec = cache.scratch.ᶜtemp_scalar_2
288+
@. ᶜprandtl_nvec = turbulent_prandtl_number(
289+
params,
290+
ᶜlinear_buoygrad,
291+
ᶜstrain_rate_norm,
292+
)
293+
ᶜmixing_length = @. lazy(
294+
smagorinsky_lilly_length(
295+
CAP.c_smag(params),
296+
sqrt(max(ᶜlinear_buoygrad, 0)), # N_eff
297+
ᶜdz,
298+
ᶜprandtl_nvec,
299+
ᶜstrain_rate_norm,
300+
),
301+
)
302+
end
303+
304+
276305
if isnothing(out)
277-
return copy(cache.precomputed.ᶜmixing_length)
306+
return copy(ᶜmixing_length)
278307
else
279-
out .= cache.precomputed.ᶜmixing_length
308+
out .= ᶜmixing_length
280309
end
281310
end,
282311
)

0 commit comments

Comments
 (0)