Skip to content

Commit 3f0631d

Browse files
costachriscostachris
committed
restructure env, specific helper functions
1 parent 41caf15 commit 3f0631d

21 files changed

+410
-310
lines changed

‎src/cache/cloud_fraction.jl

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,6 @@ NVTX.@annotate function set_cloud_fraction!(
8181
(; turbconv_model) = p.atmos
8282

8383
if isnothing(turbconv_model)
84-
ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
8584
if p.atmos.call_cloud_diagnostics_per_stage isa
8685
CallCloudDiagnosticsPerStage
8786
(; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
@@ -192,7 +191,7 @@ NVTX.@annotate function set_cloud_fraction!(
192191
(; ᶜts⁰, cloud_diagnostics_tuple) = p.precomputed
193192
(; ᶜρʲs, ᶜtsʲs) = p.precomputed
194193
(; turbconv_model) = p.atmos
195-
ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
194+
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
196195

197196
# TODO - we should make this default when using diagnostic edmf
198197
# environment
@@ -214,9 +213,9 @@ NVTX.@annotate function set_cloud_fraction!(
214213
# weight cloud diagnostics by environmental area
215214
@. cloud_diagnostics_tuple *= NamedTuple{(:cf, :q_liq, :q_ice)}(
216215
tuple(
217-
draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
218-
draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
219-
draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
216+
draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
217+
draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
218+
draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
220219
),
221220
)
222221
# updrafts

‎src/cache/diagnostic_edmf_precomputed_quantities.jl

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -105,21 +105,24 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
105105
turbconv_params = CAP.turbconv_params(params)
106106
ᶜts = p.precomputed.ᶜts #TODO replace
107107

108-
q_tot = specific(Y.c.ρq_tot, Y.c.ρ)
108+
q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
109+
ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ))
109110
ᶜh_tot = @. lazy(
110111
TD.total_specific_enthalpy(
111112
thermo_params,
112113
ᶜts,
113-
specific(Y.c.ρe_tot, Y.c.ρ),
114+
ᶜe_tot,
114115
),
115116
)
116117

117118
ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1))
118119
uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1))
119120
u³_int_halflevel = Fields.field_values(Fields.level(ᶠu³, half))
120-
h_tot_int_level = Fields.field_values(Fields.level(ᶜh_tot, 1))
121+
h_tot_int_level =
122+
Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), 1))
121123
K_int_level = Fields.field_values(Fields.level(ᶜK, 1))
122-
q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1))
124+
q_tot_int_level =
125+
Fields.field_values(Fields.level(Base.materialize(q_tot), 1))
123126

124127
p_int_level = Fields.field_values(Fields.level(ᶜp, 1))
125128
Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1))
@@ -330,7 +333,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
330333
ᶠnh_pressure³_buoyʲs,
331334
ᶠnh_pressure³_dragʲs,
332335
) = p.precomputed
333-
(; ᶠu³⁰, ᶜK⁰, ᶜtke⁰) = p.precomputed
336+
(; ᶠu³⁰, ᶜK⁰) = p.precomputed
337+
338+
334339

335340
if microphysics_model isa Microphysics1Moment
336341
q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
@@ -354,14 +359,25 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
354359
Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
355360

356361
# integral
357-
q_tot = specific(Y.c.ρq_tot, Y.c.ρ)
362+
q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
363+
ᶜh_tot = @. lazy(
364+
TD.total_specific_enthalpy(
365+
thermo_params,
366+
ᶜts,
367+
specific(Y.c.ρe_tot, Y.c.ρ),
368+
),
369+
)
370+
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
371+
358372
for i in 2:Spaces.nlevels(axes(Y.c))
359373
ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i))
360374
uₕ_level = Fields.field_values(Fields.level(Y.c.uₕ, i))
361375
u³_halflevel = Fields.field_values(Fields.level(ᶠu³, i - half))
362376
K_level = Fields.field_values(Fields.level(ᶜK, i))
363-
h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i))
364-
q_tot_level = Fields.field_values(Fields.level(q_tot, i))
377+
h_tot_level =
378+
Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), i))
379+
q_tot_level =
380+
Fields.field_values(Fields.level(Base.materialize(q_tot), i))
365381
p_level = Fields.field_values(Fields.level(ᶜp, i))
366382
Φ_level = Fields.field_values(Fields.level(ᶜΦ, i))
367383
local_geometry_level = Fields.field_values(
@@ -385,8 +401,10 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
385401
Fields.field_values(Fields.level(ᶠu³⁰, i - 1 - half))
386402
u³⁰_data_prev_halflevel = u³⁰_prev_halflevel.components.data.:1
387403
K_prev_level = Fields.field_values(Fields.level(ᶜK, i - 1))
388-
h_tot_prev_level = Fields.field_values(Fields.level(ᶜh_tot, i - 1))
389-
q_tot_prev_level = Fields.field_values(Fields.level(q_tot, i - 1))
404+
h_tot_prev_level =
405+
Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), i - 1))
406+
q_tot_prev_level =
407+
Fields.field_values(Fields.level(Base.materialize(q_tot), i - 1))
390408
ts_prev_level = Fields.field_values(Fields.level(ᶜts, i - 1))
391409
p_prev_level = Fields.field_values(Fields.level(ᶜp, i - 1))
392410
z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1))
@@ -486,7 +504,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
486504
Fields.field_values(Fields.level(ᶜq_iceʲ, i - 1))
487505
end
488506

489-
tke_prev_level = Fields.field_values(Fields.level(ᶜtke⁰, i - 1))
507+
tke_prev_level = Fields.field_values(
508+
Fields.level(Base.materialize(ᶜtke⁰), i - 1),
509+
)
490510

491511
@. entrʲ_prev_level = entrainment(
492512
thermo_params,
@@ -1049,12 +1069,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
10491069
(; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed
10501070

10511071
# Environment precipitation sources (to be applied to grid mean)
1052-
q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
1072+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
10531073
@. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(
10541074
thermo_params,
10551075
microphys_0m_params,
10561076
dt,
1057-
specific(Y.c.ρq_tot, Y.c.ρ),
1077+
ᶜq_tot,
10581078
ᶜts,
10591079
)
10601080
return nothing

‎src/cache/precipitation_precomputed_quantities.jl

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -262,14 +262,14 @@ function set_precipitation_cache!(
262262
(; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed
263263
(; ᶜts⁰, ᶜtsʲs) = p.precomputed
264264
thermo_params = CAP.thermodynamics_params(p.params)
265-
ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
265+
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
266266

267267
n = n_mass_flux_subdomains(p.atmos.turbconv_model)
268268

269-
@. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰
269+
@. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰_vals
270270
@. ᶜS_ρe_tot =
271271
ᶜSqₜᵖ⁰ *
272-
ᶜρa⁰ *
272+
ᶜρa⁰_vals *
273273
e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts⁰, ᶜΦ)
274274
for j in 1:n
275275
@. ᶜS_ρq_tot += ᶜSqₜᵖʲs.:($$j) * Y.c.sgsʲs.:($$j).ρa
@@ -289,8 +289,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _)
289289
(; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed
290290
(; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
291291

292-
q_rai = specific(Y.c.ρq_rai, Y.c.ρ)
293-
q_sno = specific(Y.c.ρq_sno, Y.c.ρ)
292+
q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
293+
q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
294+
q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ))
295+
q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
296+
q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
294297

295298
ᶜSᵖ = p.scratch.ᶜtemp_scalar
296299
ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2
@@ -364,10 +367,10 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _)
364367
(; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
365368
(; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed
366369

367-
q_liq = specific(Y.c.ρq_liq, Y.c.ρ)
368-
q_rai = specific(Y.c.ρq_rai, Y.c.ρ)
369-
n_liq = specific(Y.c.ρn_liq, Y.c.ρ)
370-
n_rai = specific(Y.c.ρn_rai, Y.c.ρ)
370+
q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
371+
q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
372+
n_liq = @. lazy(specific(Y.c.ρn_liq, Y.c.ρ))
373+
n_rai = @. lazy(specific(Y.c.ρn_rai, Y.c.ρ))
371374

372375
ᶜSᵖ = p.scratch.ᶜtemp_scalar
373376
ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2
@@ -481,20 +484,24 @@ function set_precipitation_surface_fluxes!(
481484
sfc_ρ = @. lazy(int_ρ * int_J / sfc_J)
482485

483486
# Constant extrapolation to surface, consistent with simple downwinding
487+
ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
488+
ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ))
489+
ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
490+
ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
484491
sfc_qᵣ = Fields.Field(
485-
Fields.field_values(Fields.level(specific(Y.c.ρq_rai, Y.c.ρ), 1)),
492+
Fields.field_values(Fields.level(Base.materialize(ᶜq_rai), 1)),
486493
sfc_space,
487494
)
488495
sfc_qₛ = Fields.Field(
489-
Fields.field_values(Fields.level(specific(Y.c.ρq_sno, Y.c.ρ), 1)),
496+
Fields.field_values(Fields.level(Base.materialize(ᶜq_sno), 1)),
490497
sfc_space,
491498
)
492499
sfc_qₗ = Fields.Field(
493-
Fields.field_values(Fields.level(specific(Y.c.ρq_liq, Y.c.ρ), 1)),
500+
Fields.field_values(Fields.level(Base.materialize(ᶜq_liq), 1)),
494501
sfc_space,
495502
)
496503
sfc_qᵢ = Fields.Field(
497-
Fields.field_values(Fields.level(specific(Y.c.ρq_ice, Y.c.ρ), 1)),
504+
Fields.field_values(Fields.level(Base.materialize(ᶜq_ice), 1)),
498505
sfc_space,
499506
)
500507
sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space)

‎src/cache/precomputed_quantities.jl

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -300,8 +300,9 @@ function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
300300
return nothing
301301
end
302302

303-
function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰, ᶠu₃⁰, turbconv_model)
304-
@. ᶜK += ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Yf.u₃, CT3(ᶠu₃⁰ - Yf.u₃))) / 2 / Yc.ρ
303+
function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰_vals, ᶠu₃⁰, turbconv_model)
304+
@. ᶜK +=
305+
ᶜρa⁰_vals * ᶜinterp(dot(ᶠu₃⁰ - Yf.u₃, CT3(ᶠu₃⁰ - Yf.u₃))) / 2 / Yc.ρ
305306
for j in 1:n_mass_flux_subdomains(turbconv_model)
306307
ᶜρaʲ = Y.c.sgsʲs.:($j).ρa
307308
ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃
@@ -454,8 +455,6 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
454455
set_prognostic_edmf_precomputed_quantities_draft!(Y, p, ᶠuₕ³, t)
455456
set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
456457
set_prognostic_edmf_precomputed_quantities_implicit_closures!(Y, p, t)
457-
elseif turbconv_model isa DiagnosticEDMFX
458-
set_diagnostic_edmf_precomputed_quantities!(Y, p, t)
459458
elseif !(isnothing(turbconv_model))
460459
# Do nothing for other turbconv models for now
461460
end

‎src/cache/prognostic_edmf_precomputed_quantities.jl

Lines changed: 36 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -24,22 +24,22 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
2424
(; ᶜp, ᶜK) = p.precomputed
2525
(; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed
2626

27-
ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
28-
ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))
27+
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
28+
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
2929
set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
3030
set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
3131
# @. ᶜK⁰ += ᶜtke⁰
32-
ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model))
32+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
3333

3434
ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2
3535
ᶜmse⁰ .= specific_env_mse(Y.c, p)
3636

3737
if p.atmos.moisture_model isa NonEquilMoistModel &&
3838
p.atmos.microphysics_model isa Microphysics1Moment
39-
ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model))
40-
ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model))
41-
ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model))
42-
ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model))
39+
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p)
40+
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p)
41+
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p)
42+
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p)
4343
@. ᶜts⁰ = TD.PhaseNonEquil_phq(
4444
thermo_params,
4545
ᶜp,
@@ -219,26 +219,31 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!(
219219
if p.atmos.moisture_model isa NonEquilMoistModel &&
220220
p.atmos.microphysics_model isa Microphysics1Moment
221221
# TODO - any better way to define the cloud and precip tracer flux?
222+
223+
ᶜq_liq = @.lazy(specific(Y.c.ρq_liq, Y.c.ρ))
224+
ᶜq_ice = @.lazy(specific(Y.c.ρq_ice, Y.c.ρ))
225+
ᶜq_rai = @.lazy(specific(Y.c.ρq_rai, Y.c.ρ))
226+
ᶜq_sno = @.lazy(specific(Y.c.ρq_sno, Y.c.ρ))
222227
ᶜq_liq_int_val = Fields.field_values(
223-
Fields.level(specific(Y.c.ρq_liq, Y.c.ρ), 1),
228+
Fields.level(Base.materialize(ᶜq_liq), 1),
224229
)
225230
ᶜq_liqʲ_int_val = Fields.field_values(Fields.level(ᶜq_liqʲ, 1))
226231
@. ᶜq_liqʲ_int_val = ᶜq_liq_int_val
227232

228233
ᶜq_ice_int_val = Fields.field_values(
229-
Fields.level(specific(Y.c.ρq_ice, Y.c.ρ), 1),
234+
Fields.level(Base.materialize(ᶜq_ice), 1),
230235
)
231236
ᶜq_iceʲ_int_val = Fields.field_values(Fields.level(ᶜq_iceʲ, 1))
232237
@. ᶜq_iceʲ_int_val = ᶜq_ice_int_val
233238

234239
ᶜq_rai_int_val = Fields.field_values(
235-
Fields.level(specific(Y.c.ρq_rai, Y.c.ρ), 1),
240+
Fields.level(Base.materialize(ᶜq_rai), 1),
236241
)
237242
ᶜq_raiʲ_int_val = Fields.field_values(Fields.level(ᶜq_raiʲ, 1))
238243
@. ᶜq_raiʲ_int_val = ᶜq_rai_int_val
239244

240245
ᶜq_sno_int_val = Fields.field_values(
241-
Fields.level(specific(Y.c.ρq_sno, Y.c.ρ), 1),
246+
Fields.level(Base.materialize(ᶜq_sno), 1),
242247
)
243248
ᶜq_snoʲ_int_val = Fields.field_values(Fields.level(ᶜq_snoʲ, 1))
244249
@. ᶜq_snoʲ_int_val = ᶜq_sno_int_val
@@ -360,8 +365,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
360365
ᶜdz = Fields.Δz_field(axes(Y.c))
361366
ᶜlg = Fields.local_geometry_field(Y.c)
362367
ᶠlg = Fields.local_geometry_field(Y.f)
363-
ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))
364-
ᶜρa⁰ = @. lazy(ρa⁰(Y.c))
368+
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
369+
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
365370

366371
ᶜvert_div = p.scratch.ᶜtemp_scalar
367372
ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2
@@ -446,7 +451,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
446451
(; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed
447452
# First order approximation: Use environmental mean fields.
448453
@. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))) # ∂θv∂z_unsat
449-
ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model))
454+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
450455
@. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)) # ∂qt∂z_sat
451456
@. ᶜgradᵥ_θ_liq_ice⁰ =
452457
ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))) # ∂θl∂z_sat
@@ -471,7 +476,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
471476
@. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
472477

473478
ρatke_flux_values = Fields.field_values(ρatke_flux)
474-
ρa_sfc_values = Fields.field_values(Fields.level(Base.materialize(ᶜρa⁰), 1)) # TODO: replace by surface value
479+
ρa_sfc_values =
480+
Fields.field_values(Fields.level(Base.materialize(ᶜρa⁰_vals), 1)) # TODO: replace by surface value
475481
ustar_values = Fields.field_values(ustar)
476482
sfc_local_geometry_values = Fields.field_values(
477483
Fields.level(Fields.local_geometry_field(Y.f), half),
@@ -522,7 +528,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
522528
)
523529
end
524530
# sources from the environment
525-
ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, p.atmos.turbconv_model))
531+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
526532
@. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(thp, cmp, dt, ᶜq_tot⁰, ᶜts⁰)
527533
return nothing
528534
end
@@ -612,18 +618,23 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
612618
end
613619

614620
# Precipitation sources and sinks from the environment
615-
ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model))
616-
ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model))
617-
ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model))
618-
ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model))
621+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
622+
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p)
623+
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p)
624+
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p)
625+
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p)
626+
ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰))
619627
compute_precipitation_sources!(
620628
ᶜSᵖ,
621629
ᶜSᵖ_snow,
622630
ᶜSqₗᵖ⁰,
623631
ᶜSqᵢᵖ⁰,
624632
ᶜSqᵣᵖ⁰,
625633
ᶜSqₛᵖ⁰,
626-
TD.air_density.(thp, ᶜts⁰),
634+
ᶜρ⁰,
635+
ᶜq_tot⁰,
636+
ᶜq_liq⁰,
637+
ᶜq_ice⁰,
627638
ᶜq_rai⁰,
628639
ᶜq_sno⁰,
629640
ᶜts⁰,
@@ -635,7 +646,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
635646
ᶜSᵖ,
636647
ᶜSqᵣᵖ⁰,
637648
ᶜSqₛᵖ⁰,
638-
TD.air_density.(thp, ᶜts⁰),
649+
ᶜρ⁰,
650+
ᶜq_tot⁰,
651+
ᶜq_liq⁰,
652+
ᶜq_ice⁰,
639653
ᶜq_rai⁰,
640654
ᶜq_sno⁰,
641655
ᶜts⁰,

‎src/diagnostics/Diagnostics.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ import ..ᶜcompute_eddy_diffusivity_coefficient
6767
import ..ρa⁰
6868
import ..specific_tke
6969

70+
7071
# We need the abbreviations for symbols like curl, grad, and so on
7172
include(joinpath("..", "utils", "abbreviations.jl"))
7273

0 commit comments

Comments
 (0)