Skip to content

Commit 97234db

Browse files
haakon-ehaakon-e
committed
add face density method
1 parent 576ac23 commit 97234db

File tree

6 files changed

+118
-214
lines changed

6 files changed

+118
-214
lines changed

‎src/prognostic_equations/advection.jl‎

Lines changed: 30 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
198198
advect_tke = use_prognostic_tke(turbconv_model)
199199
point_type = eltype(Fields.coordinate_field(Y.c))
200200
(; dt) = p
201-
ᶜJ = Fields.local_geometry_field(Y.c).J
201+
ᶜJ, ᶠJ = Fields.local_geometry_field(Y.c).J, Fields.local_geometry_field(Y.f).J
202+
ᶜρ, ᶠρ = Y.c.ρ, face_density(Y.c.ρ)
203+
ᶠρJ = @. lazy(ᶠρ * ᶠJ)
204+
ᶜρJ = @. lazy(ᶜρ * ᶜJ)
202205
(; ᶜf³, ᶠf¹², ᶜΦ) = p.core
203206
(; ᶜu, ᶠu³, ᶜK, ᶜts) = p.precomputed
204207
(; edmfx_upwinding) = n > 0 || advect_tke ? p.atmos.numerics : all_nothing
@@ -216,21 +219,21 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
216219
advect_tke ?
217220
(
218221
turbconv_model isa PrognosticEDMFX ?
219-
(@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
222+
(@. lazy(ρa⁰(ᶜρ, Y.c.sgsʲs, turbconv_model))) : ᶜρ
220223
) : nothing
221224
ᶜρ⁰ = if advect_tke
222225
if n > 0
223226
(; ᶜts⁰) = p.precomputed
224227
@. lazy(TD.air_density(thermo_params, ᶜts⁰))
225228
else
226-
Y.c.ρ
229+
ᶜρ
227230
end
228231
else
229232
nothing
230233
end
231234
ᶜtke⁰ =
232235
advect_tke ?
233-
(@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))) :
236+
(@. lazy(specific_tke(ᶜρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))) :
234237
nothing
235238
ᶜa_scalar = p.scratch.ᶜtemp_scalar
236239
ᶜω³ = p.scratch.ᶜtemp_CT3
@@ -253,14 +256,12 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
253256
end
254257
# Without the CT12(), the right-hand side would be a CT1 or CT2 in 2D space.
255258

256-
ᶜρ = Y.c.ρ
257-
258259
# Full vertical advection of passive tracers (like liq, rai, etc) ...
259260
# If sgs_mass_flux is true, the advection term is computed from the sum of SGS fluxes
260261
if p.atmos.edmfx_model.sgs_mass_flux isa Val{false}
261262
foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name
262263
if !(ρχ_name in (@name(ρe_tot), @name(ρq_tot)))
263-
ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
264+
ᶜχ = @. lazy(specific(ᶜρχ, ᶜρ))
264265
vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
265266
@. ᶜρχₜ += vtt
266267
end
@@ -282,17 +283,15 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
282283
end
283284

284285
if !(p.atmos.moisture_model isa DryModel) && tracer_upwinding != Val(:none)
285-
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
286+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, ᶜρ))
286287
vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), tracer_upwinding)
287288
vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), Val(:none))
288289
@. Yₜ.c.ρq_tot += vtt - vtt_central
289290
end
290291

291292
if isnothing(ᶠf¹²)
292293
# shallow atmosphere
293-
@. Yₜ.c.uₕ -=
294-
ᶜinterp(ᶠω¹² × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) / (Y.c.ρ * ᶜJ) +
295-
(ᶜf³ + ᶜω³) × CT12(ᶜu)
294+
@. Yₜ.c.uₕ -= ᶜinterp(ᶠω¹² × (ᶠρJ * ᶠu³)) / ᶜρJ + (ᶜf³ + ᶜω³) × CT12(ᶜu)
296295
@. Yₜ.f.u₃ -= ᶠω¹² × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
297296
for j in 1:n
298297
@. Yₜ.f.sgsʲs.:($$j).u₃ -=
@@ -301,9 +300,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
301300
end
302301
else
303302
# deep atmosphere
304-
@. Yₜ.c.uₕ -=
305-
ᶜinterp((ᶠf¹² + ᶠω¹²) × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) /
306-
(Y.c.ρ * ᶜJ) + (ᶜf³ + ᶜω³) × CT12(ᶜu)
303+
@. Yₜ.c.uₕ -= ᶜinterp((ᶠf¹² + ᶠω¹²) × (ᶠρJ * ᶠu³)) / ᶜρJ + (ᶜf³ + ᶜω³) × CT12(ᶜu)
307304
@. Yₜ.f.u₃ -= (ᶠf¹² + ᶠω¹²) × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
308305
for j in 1:n
309306
@. Yₜ.f.sgsʲs.:($$j).u₃ -=
@@ -348,11 +345,7 @@ Modifies EDMF updraft fields in `Yₜ.c.sgsʲs` and `Yₜ.f.sgsʲs`.
348345
"""
349346
edmfx_sgs_vertical_advection_tendency!(Yₜ, Y, p, t, turbconv_model) = nothing
350347

351-
function edmfx_sgs_vertical_advection_tendency!(
352-
Yₜ,
353-
Y,
354-
p,
355-
t,
348+
function edmfx_sgs_vertical_advection_tendency!(Yₜ, Y, p, t,
356349
turbconv_model::PrognosticEDMFX,
357350
)
358351
(; params) = p
@@ -393,28 +386,21 @@ function edmfx_sgs_vertical_advection_tendency!(
393386
end
394387

395388
for j in 1:n
396-
ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))))
389+
ᶜρʲ = ᶜρʲs.:($j)
390+
ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲ)))
397391
ᶜright_biased_∂a∂z =
398392
@. lazy(ᶜprecipdivᵥ(ᶠinterp(ᶜJ) / ᶠJ * ᶠright_bias(Geometry.WVector(ᶜa))))
399393

400394
# Flux form vertical advection of area farction with the grid mean velocity
401-
ᶜ∂ρ∂t = vertical_transport(ᶜρʲs.:($j), ᶠu³ʲs.:($j), ᶜa, dt, edmfx_upwinding)
395+
ᶜ∂ρ∂t = vertical_transport(ᶜρʲ, ᶠu³ʲs.:($j), ᶜa, dt, edmfx_upwinding)
402396
@. Yₜ.c.sgsʲs.:($$j).ρa += ᶜ∂ρ∂t
403397

404398
# Advective form advection of mse and q_tot with the grid mean velocity
405399
# Note: This allocates because the function is too long
406-
va = vertical_advection(
407-
ᶠu³ʲs.:($j),
408-
Y.c.sgsʲs.:($j).mse,
409-
edmfx_upwinding,
410-
)
400+
va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).mse, edmfx_upwinding)
411401
@. Yₜ.c.sgsʲs.:($$j).mse += va
412402

413-
va = vertical_advection(
414-
ᶠu³ʲs.:($j),
415-
Y.c.sgsʲs.:($j).q_tot,
416-
edmfx_upwinding,
417-
)
403+
va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_tot, edmfx_upwinding)
418404
@. Yₜ.c.sgsʲs.:($$j).q_tot += va
419405

420406
if p.atmos.moisture_model isa NonEquilMoistModel && (
@@ -433,15 +419,8 @@ function edmfx_sgs_vertical_advection_tendency!(
433419
ᶜ∂ρ∂t_sed = p.scratch.ᶜtemp_scalar_3
434420
@. ᶜ∂ρ∂t_sed = 0
435421

436-
ᶜinv_ρ̂ = (@. lazy(
437-
specific(
438-
FT(1),
439-
Y.c.sgsʲs.:($$j).ρa,
440-
FT(0),
441-
ᶜρʲs.:($$j),
442-
turbconv_model,
443-
),
444-
))
422+
ᶜinv_ρ̂ =
423+
@. lazy(specific(FT(1), Y.c.sgsʲs.:($$j).ρa, FT(0), ᶜρʲ, turbconv_model))
445424

446425
# Sedimentation
447426
# TODO - lazify ᶜwₗʲs computation. No need to cache it.
@@ -463,28 +442,12 @@ function edmfx_sgs_vertical_advection_tendency!(
463442
ᶜaqʲ = (@. lazy(ᶜa * ᶜqʲ))
464443

465444
# Flux form advection of tracers with updraft velocity
466-
vtt = vertical_transport(
467-
ᶜρʲs.:($j),
468-
ᶠu³ʲs.:($j),
469-
ᶜaqʲ,
470-
dt,
471-
tracer_upwinding,
472-
)
445+
vtt = vertical_transport(ᶜρʲ, ᶠu³ʲs.:($j), ᶜaqʲ, dt, tracer_upwinding)
473446
@. ᶜqʲₜ += ᶜinv_ρ̂ * (vtt - ᶜqʲ * ᶜ∂ρ∂t)
474447

475448
# Flux form sedimentation of tracers
476-
vtt = vertical_transport_sedimentation(
477-
ᶜρʲs.:($j),
478-
ᶜwʲ,
479-
ᶜaqʲ,
480-
ᶠJ,
481-
)
482-
sed_detr = sedimentation_detrainment(
483-
ᶜρʲs.:($j),
484-
ᶜwʲ,
485-
ᶜqʲ,
486-
ᶜright_biased_∂a∂z,
487-
)
449+
vtt = vertical_transport_sedimentation(ᶜρʲ, ᶜwʲ, ᶜaqʲ)
450+
sed_detr = sedimentation_detrainment(ᶜρʲ, ᶜwʲ, ᶜqʲ, ᶜright_biased_∂a∂z)
488451
@. ᶜqʲₜ += ᶜinv_ρ̂ * (vtt + sed_detr)
489452
@. Yₜ.c.sgsʲs.:($$j).q_tot += ᶜinv_ρ̂ * (vtt + sed_detr)
490453
@. ᶜ∂ρ∂t_sed += (vtt + sed_detr)
@@ -501,18 +464,11 @@ function edmfx_sgs_vertical_advection_tendency!(
501464
else
502465
error("Unsupported moisture tracer variable")
503466
end
504-
vtt = vertical_transport_sedimentation(
505-
ᶜρʲs.:($j),
506-
ᶜwʲ,
507-
ᶜaqʲ .* ᶜmse_li,
508-
ᶠJ,
509-
)
510-
sed_detr = sedimentation_detrainment(
511-
ᶜρʲs.:($j),
512-
ᶜwʲ,
513-
ᶜqʲ .* ᶜmse_li,
514-
ᶜright_biased_∂a∂z,
515-
)
467+
ᶜaqʲmse_li = @. lazy(ᶜaqʲ * ᶜmse_li)
468+
vtt = vertical_transport_sedimentation(ᶜρʲ, ᶜwʲ, ᶜaqʲmse_li)
469+
ᶜqʲmse_li = @. lazy(ᶜqʲ * ᶜmse_li)
470+
sed_detr =
471+
sedimentation_detrainment(ᶜρʲ, ᶜwʲ, ᶜqʲmse_li, ᶜright_biased_∂a∂z)
516472
@. Yₜ.c.sgsʲs.:($$j).mse += ᶜinv_ρ̂ * (vtt + sed_detr)
517473
end
518474

@@ -555,28 +511,12 @@ function edmfx_sgs_vertical_advection_tendency!(
555511
ᶜaχʲ = (@. lazy(ᶜa * ᶜχʲ))
556512

557513
# Flux form advection of tracers with updraft velocity
558-
vtt = vertical_transport(
559-
ᶜρʲs.:($j),
560-
ᶠu³ʲs.:($j),
561-
ᶜaχʲ,
562-
dt,
563-
tracer_upwinding,
564-
)
514+
vtt = vertical_transport(ᶜρʲ, ᶠu³ʲs.:($j), ᶜaχʲ, dt, tracer_upwinding)
565515
@. ᶜχʲₜ += ᶜinv_ρ̂ * (vtt - ᶜχʲ * ᶜ∂ρ∂t)
566516

567517
# Flux form sedimentation of tracers
568-
vtt = vertical_transport_sedimentation(
569-
ᶜρʲs.:($j),
570-
ᶜwʲ,
571-
ᶜaχʲ,
572-
ᶠJ,
573-
)
574-
sed_detr = sedimentation_detrainment(
575-
ᶜρʲs.:($j),
576-
ᶜwʲ,
577-
ᶜχʲ,
578-
ᶜright_biased_∂a∂z,
579-
)
518+
vtt = vertical_transport_sedimentation(ᶜρʲ, ᶜwʲ, ᶜaχʲ)
519+
sed_detr = sedimentation_detrainment(ᶜρʲ, ᶜwʲ, ᶜχʲ, ᶜright_biased_∂a∂z)
580520
@. ᶜχʲₜ += ᶜinv_ρ̂ * (vtt + sed_detr)
581521

582522
# Contribution of density variation due to sedimentation

0 commit comments

Comments
 (0)