Skip to content

Commit 9ffce7f

Browse files
haakon-ehaakon-e
committed
add face density method
1 parent 0e277a7 commit 9ffce7f

File tree

5 files changed

+102
-204
lines changed

5 files changed

+102
-204
lines changed

‎src/prognostic_equations/advection.jl‎

Lines changed: 28 additions & 92 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,13 +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,
356-
turbconv_model::PrognosticEDMFX,
357-
)
348+
function edmfx_sgs_vertical_advection_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMFX)
358349
(; params) = p
359350
n = n_prognostic_mass_flux_subdomains(turbconv_model)
360351
(; dt) = p
@@ -393,28 +384,21 @@ function edmfx_sgs_vertical_advection_tendency!(
393384
end
394385

395386
for j in 1:n
396-
ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))))
387+
ᶜρʲ = ᶜρʲs.:($j)
388+
ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲ)))
397389
ᶜright_biased_∂a∂z =
398390
@. lazy(ᶜprecipdivᵥ(ᶠinterp(ᶜJ) / ᶠJ * ᶠright_bias(Geometry.WVector(ᶜa))))
399391

400392
# 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)
393+
ᶜ∂ρ∂t = vertical_transport(ᶜρʲ, ᶠu³ʲs.:($j), ᶜa, dt, edmfx_upwinding)
402394
@. Yₜ.c.sgsʲs.:($$j).ρa += ᶜ∂ρ∂t
403395

404396
# Advective form advection of mse and q_tot with the grid mean velocity
405397
# 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-
)
398+
va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).mse, edmfx_upwinding)
411399
@. Yₜ.c.sgsʲs.:($$j).mse += va
412400

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

420404
if p.atmos.moisture_model isa NonEquilMoistModel && (
@@ -433,15 +417,7 @@ function edmfx_sgs_vertical_advection_tendency!(
433417
ᶜ∂ρ∂t_sed = p.scratch.ᶜtemp_scalar_3
434418
@. ᶜ∂ρ∂t_sed = 0
435419

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-
))
420+
ᶜinv_ρ̂ = @. lazy(specific(FT(1), Y.c.sgsʲs.:($$j).ρa, FT(0), ᶜρʲ, turbconv_model))
445421

446422
# Sedimentation
447423
# TODO - lazify ᶜwₗʲs computation. No need to cache it.
@@ -463,28 +439,12 @@ function edmfx_sgs_vertical_advection_tendency!(
463439
ᶜaqʲ = (@. lazy(ᶜa * ᶜqʲ))
464440

465441
# 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-
)
442+
vtt = vertical_transport(ᶜρʲ, ᶠu³ʲs.:($j), ᶜaqʲ, dt, tracer_upwinding)
473443
@. ᶜqʲₜ += ᶜinv_ρ̂ * (vtt - ᶜqʲ * ᶜ∂ρ∂t)
474444

475445
# 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-
)
446+
vtt = vertical_transport_sedimentation(ᶜρʲ, ᶜwʲ, ᶜaqʲ)
447+
sed_detr = sedimentation_detrainment(ᶜρʲ, ᶜwʲ, ᶜqʲ, ᶜright_biased_∂a∂z)
488448
@. ᶜqʲₜ += ᶜinv_ρ̂ * (vtt + sed_detr)
489449
@. Yₜ.c.sgsʲs.:($$j).q_tot += ᶜinv_ρ̂ * (vtt + sed_detr)
490450
@. ᶜ∂ρ∂t_sed += (vtt + sed_detr)
@@ -501,18 +461,10 @@ function edmfx_sgs_vertical_advection_tendency!(
501461
else
502462
error("Unsupported moisture tracer variable")
503463
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-
)
464+
ᶜaqʲmse_li = @. lazy(ᶜaqʲ * ᶜmse_li)
465+
vtt = vertical_transport_sedimentation(ᶜρʲ, ᶜwʲ, ᶜaqʲmse_li)
466+
ᶜqʲmse_li = @. lazy(ᶜqʲ * ᶜmse_li)
467+
sed_detr = sedimentation_detrainment(ᶜρʲ, ᶜwʲ, ᶜqʲmse_li, ᶜright_biased_∂a∂z)
516468
@. Yₜ.c.sgsʲs.:($$j).mse += ᶜinv_ρ̂ * (vtt + sed_detr)
517469
end
518470

@@ -555,28 +507,12 @@ function edmfx_sgs_vertical_advection_tendency!(
555507
ᶜaχʲ = (@. lazy(ᶜa * ᶜχʲ))
556508

557509
# 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-
)
510+
vtt = vertical_transport(ᶜρʲ, ᶠu³ʲs.:($j), ᶜaχʲ, dt, tracer_upwinding)
565511
@. ᶜχʲₜ += ᶜinv_ρ̂ * (vtt - ᶜχʲ * ᶜ∂ρ∂t)
566512

567513
# 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-
)
514+
vtt = vertical_transport_sedimentation(ᶜρʲ, ᶜwʲ, ᶜaχʲ)
515+
sed_detr = sedimentation_detrainment(ᶜρʲ, ᶜwʲ, ᶜχʲ, ᶜright_biased_∂a∂z)
580516
@. ᶜχʲₜ += ᶜinv_ρ̂ * (vtt + sed_detr)
581517

582518
# Contribution of density variation due to sedimentation

0 commit comments

Comments
 (0)