Skip to content

Commit 34d9fa8

Browse files
haakon-ehaakon-e
committed
rft: clean up hyperdiffusion
1 parent 1e9bbb8 commit 34d9fa8

File tree

1 file changed

+44
-86
lines changed

1 file changed

+44
-86
lines changed

‎src/prognostic_equations/hyperdiffusion.jl‎

Lines changed: 44 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -22,26 +22,16 @@ function ν₄(hyperdiff, Y)
2222
return (; ν₄_scalar, ν₄_vorticity)
2323
end
2424

25-
hyperdiffusion_cache(Y, atmos) = hyperdiffusion_cache(
26-
Y,
27-
atmos.hyperdiff,
28-
atmos.turbconv_model,
29-
atmos.moisture_model,
30-
atmos.microphysics_model,
31-
)
32-
33-
# No hyperdiffiusion
34-
hyperdiffusion_cache(Y, hyperdiff::Nothing, _, _, _) = (;)
25+
function hyperdiffusion_cache(Y, atmos)
26+
(; hyperdiff, turbconv_model, moisture_model, microphysics_model) = atmos
27+
isnothing(hyperdiff) && return (;) # No hyperdiffiusion
28+
hyperdiffusion_cache(Y, hyperdiff, turbconv_model, moisture_model, microphysics_model)
29+
end
3530

3631
function hyperdiffusion_cache(
37-
Y,
38-
hyperdiff::ClimaHyperdiffusion,
39-
turbconv_model,
40-
moisture_model,
41-
microphysics_model,
32+
Y, hyperdiff::ClimaHyperdiffusion, turbconv_model, moisture_model, microphysics_model,
4233
)
43-
quadrature_style =
44-
Spaces.quadrature_style(Spaces.horizontal_space(axes(Y.c)))
34+
quadrature_style = Spaces.quadrature_style(Spaces.horizontal_space(axes(Y.c)))
4535
FT = eltype(Y)
4636
n = n_mass_flux_subdomains(turbconv_model)
4737

@@ -54,9 +44,7 @@ function hyperdiffusion_cache(
5444
)
5545

5646
# Sub-grid scale quantities
57-
ᶜ∇²uʲs =
58-
turbconv_model isa PrognosticEDMFX ? similar(Y.c, NTuple{n, C123{FT}}) :
59-
(;)
47+
ᶜ∇²uʲs = turbconv_model isa PrognosticEDMFX ? similar(Y.c, NTuple{n, C123{FT}}) : (;)
6048
moisture_sgs_quantities =
6149
moisture_model isa NonEquilMoistModel &&
6250
microphysics_model isa Microphysics1Moment ?
@@ -86,17 +74,13 @@ function hyperdiffusion_cache(
8674
moisture_sgs_quantities...,
8775
) : (;)
8876
maybe_ᶜ∇²tke⁰ =
89-
use_prognostic_tke(turbconv_model) ? (; ᶜ∇²tke⁰ = similar(Y.c, FT)) :
90-
(;)
77+
use_prognostic_tke(turbconv_model) ? (; ᶜ∇²tke⁰ = similar(Y.c, FT)) : (;)
9178
sgs_quantities = (; sgs_quantities..., maybe_ᶜ∇²tke⁰...)
9279
quantities = (; gs_quantities..., sgs_quantities...)
9380
if do_dss(axes(Y.c))
9481
quantities = (;
9582
quantities...,
96-
hyperdiffusion_ghost_buffer = map(
97-
Spaces.create_dss_buffer,
98-
quantities,
99-
),
83+
hyperdiffusion_ghost_buffer = map(Spaces.create_dss_buffer, quantities),
10084
)
10185
end
10286
return (; quantities..., ᶜ∇²u, ᶜ∇²uʲs)
@@ -115,28 +99,25 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)
11599

116100
n = n_mass_flux_subdomains(turbconv_model)
117101
diffuse_tke = use_prognostic_tke(turbconv_model)
118-
(; ᶜp) = p.precomputed
102+
(; ᶜp, ᶜu) = p.precomputed
119103
(; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff
120104
if turbconv_model isa PrognosticEDMFX
121105
(; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
106+
(; ᶜuʲs) = p.precomputed
122107
end
123108

124109
# Grid scale hyperdiffusion
125-
@. ᶜ∇²u =
126-
C123(wgradₕ(divₕ(p.precomputed.ᶜu))) -
127-
C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜu))))
110+
@. ᶜ∇²u = C123(wgradₕ(divₕ(ᶜu))) - C123(wcurlₕ(C123(curlₕ(ᶜu))))
128111

129112
ᶜh_ref = @. lazy(h_dr(thermo_params, ᶜts, ᶜΦ))
130113

131-
@. ᶜ∇²specific_energy =
132-
wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref))
114+
@. ᶜ∇²specific_energy = wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref))
133115

134116
if diffuse_tke
135117
ᶜρa⁰ =
136118
turbconv_model isa PrognosticEDMFX ?
137119
(@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
138-
ᶜtke⁰ =
139-
@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
120+
ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
140121
(; ᶜ∇²tke⁰) = p.hyperdiff
141122
@. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰))
142123
end
@@ -145,8 +126,7 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)
145126
if turbconv_model isa PrognosticEDMFX
146127
for j in 1:n
147128
@. ᶜ∇²uʲs.:($$j) =
148-
C123(wgradₕ(divₕ(p.precomputed.ᶜuʲs.:($$j)))) -
149-
C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜuʲs.:($$j)))))
129+
C123(wgradₕ(divₕ(ᶜuʲs.:($$j)))) - C123(wcurlₕ(C123(curlₕ(ᶜuʲs.:($$j)))))
150130
@. ᶜ∇²mseʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).mse))
151131
@. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j))
152132
@. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j))
@@ -183,11 +163,11 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
183163
end
184164

185165
# re-use to store the curl-curl part
186-
@. ᶜ∇²u =
166+
ᶜ∇⁴u = @. ᶜ∇²u =
187167
divergence_damping_factor * C123(wgradₕ(divₕ(ᶜ∇²u))) -
188168
C123(wcurlₕ(C123(curlₕ(ᶜ∇²u))))
189-
@. Yₜ.c.uₕ -= ν₄_vorticity * C12(ᶜ∇²u)
190-
@. Yₜ.f.u₃ -= ν₄_vorticity * ᶠwinterp(ᶜJ * Y.c.ρ, C3(ᶜ∇²u))
169+
@. Yₜ.c.uₕ -= ν₄_vorticity * C12(ᶜ∇⁴u)
170+
@. Yₜ.f.u₃ -= ν₄_vorticity * ᶠwinterp(ᶜJ * Y.c.ρ, C3(ᶜ∇⁴u))
191171

192172
@. Yₜ.c.ρe_tot -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²specific_energy))
193173

@@ -199,13 +179,11 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
199179
for j in 1:n
200180
if point_type <: Geometry.Abstract3DPoint
201181
# only need curl-curl part
202-
@. ᶜ∇²uᵥʲs.:($$j) = C3(wcurlₕ(C123(curlₕ(ᶜ∇²uʲs.:($$j)))))
203-
@. Yₜ.f.sgsʲs.:($$j).u₃ +=
204-
ν₄_vorticity * ᶠwinterp(ᶜJ * Y.c.ρ, ᶜ∇²uᵥʲs.:($$j))
182+
ᶜ∇⁴uᵥʲs = @. ᶜ∇²uᵥʲs.:($$j) = C3(wcurlₕ(C123(curlₕ(ᶜ∇²uʲs.:($$j)))))
183+
@. Yₜ.f.sgsʲs.:($$j).u₃ += ν₄_vorticity * ᶠwinterp(ᶜJ * Y.c.ρ, ᶜ∇⁴uᵥʲs)
205184
end
206185
# Note: It is more correct to have ρa inside and outside the divergence
207-
@. Yₜ.c.sgsʲs.:($$j).mse -=
208-
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²mseʲs.:($$j)))
186+
@. Yₜ.c.sgsʲs.:($$j).mse -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²mseʲs.:($$j)))
209187
end
210188
end
211189

@@ -268,15 +246,14 @@ function dss_hyperdiffusion_tendency_pairs(p)
268246
p.hyperdiff.ᶜ∇²q_liqʲs => buffer.ᶜ∇²q_liqʲs,
269247
p.hyperdiff.ᶜ∇²q_raiʲs => buffer.ᶜ∇²q_raiʲs,
270248
) : ()
271-
tracer_pairs =
272-
(core_tracer_pairs..., tc_tracer_pairs..., tc_moisture_pairs...)
249+
tracer_pairs = (core_tracer_pairs..., tc_tracer_pairs..., tc_moisture_pairs...)
273250
return (dynamics_pairs..., tracer_pairs...)
274251
end
275252

276253
# This should prep variables that we will dss in
277254
# dss_hyperdiffusion_tendency_pairs
278255
NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
279-
(; hyperdiff, turbconv_model) = p.atmos
256+
(; hyperdiff, turbconv_model, moisture_model, microphysics_model) = p.atmos
280257
isnothing(hyperdiff) && return nothing
281258

282259
(; ᶜ∇²specific_tracers) = p.hyperdiff
@@ -294,8 +271,8 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
294271
# Note: It is more correct to have ρa inside and outside the divergence
295272
@. ᶜ∇²q_totʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_tot))
296273
end
297-
if p.atmos.moisture_model isa NonEquilMoistModel &&
298-
p.atmos.microphysics_model isa Microphysics1Moment
274+
if moisture_model isa NonEquilMoistModel &&
275+
microphysics_model isa Microphysics1Moment
299276
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs) = p.hyperdiff
300277
for j in 1:n
301278
# Note: It is more correct to have ρa inside and outside the divergence
@@ -304,16 +281,10 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
304281
@. ᶜ∇²q_raiʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_rai))
305282
@. ᶜ∇²q_snoʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_sno))
306283
end
307-
elseif p.atmos.moisture_model isa NonEquilMoistModel &&
308-
p.atmos.microphysics_model isa Microphysics2Moment
309-
(;
310-
ᶜ∇²q_liqʲs,
311-
ᶜ∇²q_iceʲs,
312-
ᶜ∇²q_raiʲs,
313-
ᶜ∇²q_snoʲs,
314-
ᶜ∇²n_liqʲs,
315-
ᶜ∇²n_raiʲs,
316-
) = p.hyperdiff
284+
elseif moisture_model isa NonEquilMoistModel &&
285+
microphysics_model isa Microphysics2Moment
286+
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs, ᶜ∇²n_liqʲs, ᶜ∇²n_raiʲs) =
287+
p.hyperdiff
317288
for j in 1:n
318289
# Note: It is more correct to have ρa inside and outside the divergence
319290
@. ᶜ∇²q_liqʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_liq))
@@ -331,7 +302,7 @@ end
331302
# This requires dss to have been called on
332303
# variables in dss_hyperdiffusion_tendency_pairs
333304
NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
334-
(; hyperdiff, turbconv_model) = p.atmos
305+
(; hyperdiff, turbconv_model, moisture_model, microphysics_model) = p.atmos
335306
isnothing(hyperdiff) && return nothing
336307

337308
# Rescale the hyperdiffusivity for precipitating species.
@@ -359,41 +330,28 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
359330
(; ᶜ∇²q_totʲs) = p.hyperdiff
360331
for j in 1:n
361332
@. Yₜ.c.sgsʲs.:($$j).ρa -=
362-
ν₄_scalar *
363-
wdivₕ(Y.c.sgsʲs.:($$j).ρa * gradₕ(ᶜ∇²q_totʲs.:($$j)))
364-
@. Yₜ.c.sgsʲs.:($$j).q_tot -=
365-
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_totʲs.:($$j)))
333+
ν₄_scalar * wdivₕ(Y.c.sgsʲs.:($$j).ρa * gradₕ(ᶜ∇²q_totʲs.:($$j)))
334+
@. Yₜ.c.sgsʲs.:($$j).q_tot -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_totʲs.:($$j)))
366335
end
367-
if p.atmos.moisture_model isa NonEquilMoistModel &&
368-
p.atmos.microphysics_model isa Microphysics1Moment
336+
if moisture_model isa NonEquilMoistModel &&
337+
microphysics_model isa Microphysics1Moment
369338
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs) = p.hyperdiff
370339
for j in 1:n
371-
@. Yₜ.c.sgsʲs.:($$j).q_liq -=
372-
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
373-
@. Yₜ.c.sgsʲs.:($$j).q_ice -=
374-
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
340+
@. Yₜ.c.sgsʲs.:($$j).q_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
341+
@. Yₜ.c.sgsʲs.:($$j).q_ice -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
375342
@. Yₜ.c.sgsʲs.:($$j).q_rai -=
376343
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j)))
377344
@. Yₜ.c.sgsʲs.:($$j).q_sno -=
378345
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_snoʲs.:($$j)))
379346
end
380-
elseif p.atmos.moisture_model isa NonEquilMoistModel &&
381-
p.atmos.microphysics_model isa Microphysics2Moment
382-
(;
383-
ᶜ∇²q_liqʲs,
384-
ᶜ∇²q_iceʲs,
385-
ᶜ∇²q_raiʲs,
386-
ᶜ∇²q_snoʲs,
387-
ᶜ∇²n_liqʲs,
388-
ᶜ∇²n_raiʲs,
389-
) = p.hyperdiff
347+
elseif moisture_model isa NonEquilMoistModel &&
348+
microphysics_model isa Microphysics2Moment
349+
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs, ᶜ∇²n_liqʲs, ᶜ∇²n_raiʲs) =
350+
p.hyperdiff
390351
for j in 1:n
391-
@. Yₜ.c.sgsʲs.:($$j).q_liq -=
392-
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
393-
@. Yₜ.c.sgsʲs.:($$j).q_ice -=
394-
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
395-
@. Yₜ.c.sgsʲs.:($$j).n_liq -=
396-
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²n_liqʲs.:($$j)))
352+
@. Yₜ.c.sgsʲs.:($$j).q_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
353+
@. Yₜ.c.sgsʲs.:($$j).q_ice -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
354+
@. Yₜ.c.sgsʲs.:($$j).n_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²n_liqʲs.:($$j)))
397355
@. Yₜ.c.sgsʲs.:($$j).q_rai -=
398356
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j)))
399357
@. Yₜ.c.sgsʲs.:($$j).q_sno -=

0 commit comments

Comments
 (0)