Skip to content

Commit 712559e

Browse files
trontryteltrontrytel
authored
Merge pull request #3905 from CliMA/aj/diag_edmf_refactor
Small refactor
2 parents 477c58c + 5b30999 commit 712559e

File tree

1 file changed

+154
-153
lines changed

1 file changed

+154
-153
lines changed

‎src/cache/diagnostic_edmf_precomputed_quantities.jl‎

Lines changed: 154 additions & 153 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,46 @@ import NVTX
55
import Thermodynamics as TD
66
import ClimaCore: Spaces, Fields, RecursiveApply
77

8+
###
9+
### Helper functions for the diagnosic edmf integral
10+
###
11+
# ᶠJ - Jacobian at half level below (-1/2)
12+
# ᶜJₚ - Jacobian at previous level below (-1)
13+
# ᶠJₚ - Jacobian at previous half level below (-3/2)
14+
# ᶜρaʲₚ - updraft density * area at previous level below (-1)
15+
# ᶜρʲₚ - updraft density at previous level below (-1)
16+
# ᶜρₚ - environment density at previous level below (-1)
17+
# ᶜ∇ϕ₃ₚ - covariant geopotential gradient at previous level below (-1)
18+
# ᶠu³ʲₚ - contravariant updraft velocity at previous half level below [1/s] (-3/2)
19+
# ᶜϵʲₚ - entrainment at previous level (-1)
20+
# ᶜδʲₚ - detrainment at previous level (-1)
21+
# ᶜϵₜʲₚ - turbulent entrainment at previous level (-1)
22+
# ᶜSʲₚ - microphysics sources and sinks at previous level (-1)
23+
# ᶜtracerʲₚ - updraft property at previous level (-1)
24+
# ᶜtracerₚ - environment property at previous level (-1)
25+
26+
# Advection of area, mse and tracers
27+
function diag_edmf_advection(ᶠJ, ᶠJₚ, ᶜρaʲₚ, ᶠu³ʲₚ, ᶜtracerʲₚ)
28+
return (1 / ᶠJ) * (ᶠJₚ * ᶜρaʲₚ * ᶠu³ʲₚ * ᶜtracerʲₚ)
29+
end
30+
# Entrainment/detrainment of area, mse and tracers
31+
# Note that updraft area entrainment does not include turbulent entrainment.
32+
# In order to re-use the same function for all tracers, we pass in ones
33+
# as updraft and environment tracers for area fraction.
34+
function entr_detr(ᶠJ, ᶜJₚ, ᶜρaʲₚ, ᶜϵʲₚ, ᶜδʲₚ, ᶜϵₜʲₚ, ᶜtracerₚ, ᶜtracerʲₚ)
35+
return (1 / ᶠJ) * (
36+
ᶜJₚ * ᶜρaʲₚ * ((ᶜϵʲₚ + ᶜϵₜʲₚ) * ᶜtracerₚ - (ᶜδʲₚ + ᶜϵₜʲₚ) * ᶜtracerʲₚ)
37+
)
38+
end
39+
# Buoyancy term for mse
40+
function mse_buoyancy(ᶠJ, ᶜJₚ, ᶜρaʲₚ, ᶠu³ʲₚ, ᶜρʲₚ, ᶜρₚ, ᶜ∇Φ₃ₚ)
41+
return (1 / ᶠJ) * (ᶜJₚ * ᶜρaʲₚ * ᶠu³ʲₚ * (ᶜρʲₚ - ᶜρₚ) / ᶜρʲₚ * ᶜ∇Φ₃ₚ)
42+
end
43+
# Microphysics sources
44+
function microphysics_sources(ᶠJ, ᶜJₚ, ᶜρaʲₚ, ᶜSʲₚ)
45+
return (1 / ᶠJ) * (ᶜJₚ * ᶜρaʲₚ * ᶜSʲₚ)
46+
end
47+
848
@inline function kinetic_energy(
949
uₕ_level,
1050
u³_halflevel,
@@ -268,33 +308,6 @@ function compute_u³ʲ_u³ʲ(
268308
return u³ʲ_u³ʲ
269309
end
270310

271-
function compute_ρaʲu³ʲ(
272-
J_halflevel,
273-
J_prev_level,
274-
J_prev_halflevel,
275-
ρaʲ_prev_level,
276-
entrʲ_prev_level,
277-
detrʲ_prev_level,
278-
u³ʲ_data_prev_halflevel,
279-
S_q_totʲ_prev_level,
280-
microphysics_model,
281-
)
282-
283-
ρaʲu³ʲ_data =
284-
(1 / J_halflevel) *
285-
(J_prev_halflevel * ρaʲ_prev_level * u³ʲ_data_prev_halflevel)
286-
287-
ρaʲu³ʲ_data +=
288-
(1 / J_halflevel) *
289-
(J_prev_level * ρaʲ_prev_level * (entrʲ_prev_level - detrʲ_prev_level))
290-
if microphysics_model isa Union{Microphysics0Moment, Microphysics1Moment}
291-
ρaʲu³ʲ_data +=
292-
(1 / J_halflevel) *
293-
(J_prev_level * ρaʲ_prev_level * S_q_totʲ_prev_level)
294-
end
295-
return ρaʲu³ʲ_data
296-
end
297-
298311
NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
299312
Y,
300313
p,
@@ -696,160 +709,148 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
696709
ρaʲu³ʲ_data = p.scratch.temp_data_level_2
697710
ρaʲu³ʲ_datamse = ρaʲu³ʲ_dataq_tot = p.scratch.temp_data_level_3
698711

699-
@. ρaʲu³ʲ_data = compute_ρaʲu³ʲ(
700-
local_geometry_halflevel.J,
701-
local_geometry_prev_level.J,
702-
local_geometry_prev_halflevel.J,
703-
ρaʲ_prev_level,
704-
entrʲ_prev_level,
705-
detrʲ_prev_level,
706-
u³ʲ_data_prev_halflevel,
707-
S_q_totʲ_prev_level,
708-
microphysics_model,
709-
)
712+
###
713+
### Area fraction
714+
###
715+
@. ρaʲu³ʲ_data =
716+
diag_edmf_advection(
717+
local_geometry_halflevel.J,
718+
local_geometry_prev_halflevel.J,
719+
ρaʲ_prev_level,
720+
u³ʲ_data_prev_halflevel,
721+
FT(1),
722+
) + entr_detr(
723+
local_geometry_halflevel.J,
724+
local_geometry_prev_level.J,
725+
ρaʲ_prev_level,
726+
entrʲ_prev_level,
727+
detrʲ_prev_level,
728+
turb_entrʲ_prev_level,
729+
FT(1),
730+
FT(1),
731+
)
732+
if microphysics_model isa Microphysics0Moment
733+
@. ρaʲu³ʲ_data += microphysics_sources(
734+
local_geometry_halflevel.J,
735+
local_geometry_prev_level.J,
736+
ρaʲ_prev_level,
737+
S_q_totʲ_prev_level,
738+
)
739+
end
710740

741+
# Change current level velocity and density * area fraction
742+
kill_updraft = @. lazy(
743+
(u³ʲ_datau³ʲ_data < 10 * ∇Φ³_data_prev_level * eps(FT)) |
744+
(ρaʲu³ʲ_data < (minimum_value / ∂x³∂ξ³_level)),
745+
)
711746
@. u³ʲ_halflevel = ifelse(
712-
(
713-
(u³ʲ_datau³ʲ_data < 10 * ∇Φ³_data_prev_level * eps(FT)) | (ρaʲu³ʲ_data < (minimum_value / ∂x³∂ξ³_level))
714-
),
747+
kill_updraft,
715748
u³_halflevel,
716-
CT3(sqrt(max(0, u³ʲ_datau³ʲ_data))),
749+
CT3(sqrt(max(FT(0), u³ʲ_datau³ʲ_data))),
717750
)
718751
@. ρaʲ_level = ifelse(
719-
(
720-
(u³ʲ_datau³ʲ_data < 10 * ∇Φ³_data_prev_level * eps(FT)) | (ρaʲu³ʲ_data < (minimum_value / ∂x³∂ξ³_level))
721-
),
722-
0,
723-
ρaʲu³ʲ_data / sqrt(max(0, u³ʲ_datau³ʲ_data)),
752+
kill_updraft,
753+
FT(0),
754+
ρaʲu³ʲ_data / sqrt(max(FT(0), u³ʲ_datau³ʲ_data)),
724755
)
725756

757+
###
758+
### Moist static energy
759+
###
726760
@. ρaʲu³ʲ_datamse =
727-
(1 / local_geometry_halflevel.J) * (
728-
local_geometry_prev_halflevel.J *
729-
ρaʲ_prev_level *
730-
u³ʲ_data_prev_halflevel *
731-
mseʲ_prev_level
732-
)
733-
@. ρaʲu³ʲ_datamse +=
734-
(1 / local_geometry_halflevel.J) * (
735-
local_geometry_prev_level.J *
736-
ρaʲ_prev_level *
737-
u³ʲ_data_prev_halflevel *
738-
(ρʲ_prev_level - ρ_prev_level) / ρʲ_prev_level *
739-
∇Φ₃_data_prev_level
740-
)
741-
@. ρaʲu³ʲ_datamse +=
742-
(1 / local_geometry_halflevel.J) * (
743-
local_geometry_prev_level.J *
744-
ρaʲ_prev_level *
745-
(
746-
(entrʲ_prev_level + turb_entrʲ_prev_level) *
747-
(h_tot_prev_level - K_prev_level) -
748-
(detrʲ_prev_level + turb_entrʲ_prev_level) *
749-
mseʲ_prev_level
750-
)
761+
diag_edmf_advection(
762+
local_geometry_halflevel.J,
763+
local_geometry_prev_halflevel.J,
764+
ρaʲ_prev_level,
765+
u³ʲ_data_prev_halflevel,
766+
mseʲ_prev_level,
767+
) +
768+
mse_buoyancy(
769+
local_geometry_halflevel.J,
770+
local_geometry_prev_level.J,
771+
ρaʲ_prev_level,
772+
u³ʲ_data_prev_halflevel,
773+
ρʲ_prev_level,
774+
ρ_prev_level,
775+
∇Φ₃_data_prev_level,
776+
) +
777+
entr_detr(
778+
local_geometry_halflevel.J,
779+
local_geometry_prev_level.J,
780+
ρaʲ_prev_level,
781+
entrʲ_prev_level,
782+
detrʲ_prev_level,
783+
turb_entrʲ_prev_level,
784+
h_tot_prev_level - K_prev_level,
785+
mseʲ_prev_level,
751786
)
752787
if microphysics_model isa Microphysics0Moment
753-
@. ρaʲu³ʲ_datamse +=
754-
(1 / local_geometry_halflevel.J) * (
755-
local_geometry_prev_level.J *
756-
ρaʲ_prev_level *
757-
(
758-
S_q_totʲ_prev_level *
759-
e_tot_0M_precipitation_sources_helper(
760-
thermo_params,
761-
tsʲ_prev_level,
762-
Φ_prev_level,
763-
)
764-
)
765-
)
766-
elseif microphysics_model isa Microphysics1Moment
767-
@. ρaʲu³ʲ_datamse +=
768-
(1 / local_geometry_halflevel.J) * (
769-
local_geometry_prev_level.J *
770-
ρaʲ_prev_level *
771-
S_e_totʲ_prev_level
772-
)
788+
@. ρaʲu³ʲ_datamse += microphysics_sources(
789+
local_geometry_halflevel.J,
790+
local_geometry_prev_level.J,
791+
ρaʲ_prev_level,
792+
S_q_totʲ_prev_level *
793+
e_tot_0M_precipitation_sources_helper(
794+
thermo_params,
795+
tsʲ_prev_level,
796+
Φ_prev_level,
797+
),
798+
)
773799
end
774-
775800
@. mseʲ_level = ifelse(
776-
(
777-
(u³ʲ_datau³ʲ_data < 10 * ∇Φ³_data_prev_level * eps(FT)) | (ρaʲu³ʲ_data < (minimum_value / ∂x³∂ξ³_level))
778-
),
801+
kill_updraft,
779802
h_tot_level - K_level,
780803
ρaʲu³ʲ_datamse / ρaʲu³ʲ_data,
781804
)
782805

806+
###
807+
### Total water
808+
###
783809
@. ρaʲu³ʲ_dataq_tot =
784-
(1 / local_geometry_halflevel.J) * (
785-
local_geometry_prev_halflevel.J *
786-
ρaʲ_prev_level *
787-
u³ʲ_data_prev_halflevel *
788-
q_totʲ_prev_level
810+
diag_edmf_advection(
811+
local_geometry_halflevel.J,
812+
local_geometry_prev_halflevel.J,
813+
ρaʲ_prev_level,
814+
u³ʲ_data_prev_halflevel,
815+
q_totʲ_prev_level,
816+
) + entr_detr(
817+
local_geometry_halflevel.J,
818+
local_geometry_prev_level.J,
819+
ρaʲ_prev_level,
820+
entrʲ_prev_level,
821+
detrʲ_prev_level,
822+
turb_entrʲ_prev_level,
823+
q_tot_prev_level,
824+
q_totʲ_prev_level,
789825
)
790-
@. ρaʲu³ʲ_dataq_tot +=
791-
(1 / local_geometry_halflevel.J) * (
792-
local_geometry_prev_level.J *
793-
ρaʲ_prev_level *
794-
(
795-
(entrʲ_prev_level + turb_entrʲ_prev_level) *
796-
q_tot_prev_level -
797-
(detrʲ_prev_level + turb_entrʲ_prev_level) *
798-
q_totʲ_prev_level
799-
)
826+
if microphysics_model isa Microphysics0Moment
827+
@. ρaʲu³ʲ_dataq_tot += microphysics_sources(
828+
local_geometry_halflevel.J,
829+
local_geometry_prev_level.J,
830+
ρaʲ_prev_level,
831+
S_q_totʲ_prev_level,
800832
)
801-
if microphysics_model isa
802-
Union{Microphysics0Moment, Microphysics1Moment}
803-
@. ρaʲu³ʲ_dataq_tot +=
804-
(1 / local_geometry_halflevel.J) * (
805-
local_geometry_prev_level.J *
806-
ρaʲ_prev_level *
807-
S_q_totʲ_prev_level
808-
)
809833
end
810-
811834
@. q_totʲ_level = ifelse(
812-
(
813-
(u³ʲ_datau³ʲ_data < 10 * ∇Φ³_data_prev_level * eps(FT)) | (ρaʲu³ʲ_data < (minimum_value / ∂x³∂ξ³_level))
814-
),
835+
kill_updraft,
815836
q_tot_level,
816837
ρaʲu³ʲ_dataq_tot / ρaʲu³ʲ_data,
817838
)
818839

819840
# set updraft to grid-mean if vertical velocity is too small
820841
if i > 2
821-
@. ρaʲ_level = ifelse(
822-
(
823-
u³ʲ_data_prev_halflevel * u³ʲ_data_prev_halflevel <
824-
∇Φ³_data_prev_level * (ρʲ_prev_level - ρ_prev_level) / ρʲ_prev_level
825-
),
826-
0,
827-
ρaʲ_level,
828-
)
829-
@. u³ʲ_halflevel = ifelse(
830-
(
831-
u³ʲ_data_prev_halflevel * u³ʲ_data_prev_halflevel <
832-
∇Φ³_data_prev_level * (ρʲ_prev_level - ρ_prev_level) / ρʲ_prev_level
833-
),
834-
u³_halflevel,
835-
u³ʲ_halflevel,
836-
)
837-
@. mseʲ_level = ifelse(
838-
(
839-
u³ʲ_data_prev_halflevel * u³ʲ_data_prev_halflevel <
840-
∇Φ³_data_prev_level * (ρʲ_prev_level - ρ_prev_level) / ρʲ_prev_level
841-
),
842-
h_tot_level - K_level,
843-
mseʲ_level,
844-
)
845-
@. q_totʲ_level = ifelse(
846-
(
847-
u³ʲ_data_prev_halflevel * u³ʲ_data_prev_halflevel <
848-
∇Φ³_data_prev_level * (ρʲ_prev_level - ρ_prev_level) / ρʲ_prev_level
849-
),
850-
q_tot_level,
851-
q_totʲ_level,
842+
kill_updraft_2 = @. lazy(
843+
u³ʲ_data_prev_halflevel * u³ʲ_data_prev_halflevel <
844+
∇Φ³_data_prev_level * (ρʲ_prev_level - ρ_prev_level) /
845+
ρʲ_prev_level,
852846
)
847+
@. ρaʲ_level = ifelse(kill_updraft_2, FT(0), ρaʲ_level)
848+
@. u³ʲ_halflevel =
849+
ifelse(kill_updraft_2, u³_halflevel, u³ʲ_halflevel)
850+
@. mseʲ_level =
851+
ifelse(kill_updraft_2, h_tot_level - K_level, mseʲ_level)
852+
@. q_totʲ_level =
853+
ifelse(kill_updraft_2, q_tot_level, q_totʲ_level)
853854
end
854855

855856
@. Kʲ_level = kinetic_energy(

0 commit comments

Comments
 (0)