Skip to content

Commit a3d73be

Browse files
haakon-ehaakon-e
committed
rft: clean up hyperdiffusion
1 parent 1fd80d6 commit a3d73be

File tree

1 file changed

+48
-89
lines changed

1 file changed

+48
-89
lines changed

‎src/prognostic_equations/hyperdiffusion.jl‎

Lines changed: 48 additions & 89 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))
@@ -165,11 +145,12 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
165145

166146
n = n_mass_flux_subdomains(turbconv_model)
167147
diffuse_tke = use_prognostic_tke(turbconv_model)
148+
ᶜρ = Y.c.ρ
168149
ᶜJ = Fields.local_geometry_field(Y.c).J
169150
point_type = eltype(Fields.coordinate_field(Y.c))
170151
(; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff
171152
if turbconv_model isa PrognosticEDMFX
172-
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
153+
ᶜρa⁰ = @. lazy(ρa⁰(ᶜρ, Y.c.sgsʲs, turbconv_model))
173154
(; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
174155
end
175156
if use_prognostic_tke(turbconv_model)
@@ -183,13 +164,13 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
183164
end
184165

185166
# re-use to store the curl-curl part
186-
@. ᶜ∇²u =
167+
ᶜ∇⁴u = @. ᶜ∇²u =
187168
divergence_damping_factor * C123(wgradₕ(divₕ(ᶜ∇²u))) -
188169
C123(wcurlₕ(C123(curlₕ(ᶜ∇²u))))
189-
@. Yₜ.c.uₕ -= ν₄_vorticity * C12(ᶜ∇²u)
190-
@. Yₜ.f.u₃ -= ν₄_vorticity * ᶠwinterp(ᶜJ * Y.c.ρ, C3(ᶜ∇²u))
170+
@. Yₜ.c.uₕ -= ν₄_vorticity * C12(ᶜ∇⁴u)
171+
@. Yₜ.f.u₃ -= ν₄_vorticity * ᶠwinterp(ᶜJ * ᶜρ, C3(ᶜ∇⁴u))
191172

192-
@. Yₜ.c.ρe_tot -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²specific_energy))
173+
@. Yₜ.c.ρe_tot -= ν₄_scalar * wdivₕ(ᶜρ * gradₕ(ᶜ∇²specific_energy))
193174

194175
# Sub-grid scale hyperdiffusion continued
195176
if (turbconv_model isa PrognosticEDMFX) && diffuse_tke
@@ -199,18 +180,16 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
199180
for j in 1:n
200181
if point_type <: Geometry.Abstract3DPoint
201182
# 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))
183+
ᶜ∇⁴uᵥʲ = @. ᶜ∇²uᵥʲs.:($$j) = C3(wcurlₕ(C123(curlₕ(ᶜ∇²uʲs.:($$j)))))
184+
@. Yₜ.f.sgsʲs.:($$j).u₃ += ν₄_vorticity * ᶠwinterp(ᶜJ * ᶜρ, ᶜ∇⁴uᵥʲ)
205185
end
206186
# 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)))
187+
@. Yₜ.c.sgsʲs.:($$j).mse -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²mseʲs.:($$j)))
209188
end
210189
end
211190

212191
if turbconv_model isa DiagnosticEDMFX && diffuse_tke
213-
@. Yₜ.c.sgs⁰.ρatke -= ν₄_vorticity * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²tke⁰))
192+
@. Yₜ.c.sgs⁰.ρatke -= ν₄_vorticity * wdivₕ(ᶜρ * gradₕ(ᶜ∇²tke⁰))
214193
end
215194
end
216195

@@ -268,15 +247,14 @@ function dss_hyperdiffusion_tendency_pairs(p)
268247
p.hyperdiff.ᶜ∇²q_liqʲs => buffer.ᶜ∇²q_liqʲs,
269248
p.hyperdiff.ᶜ∇²q_raiʲs => buffer.ᶜ∇²q_raiʲs,
270249
) : ()
271-
tracer_pairs =
272-
(core_tracer_pairs..., tc_tracer_pairs..., tc_moisture_pairs...)
250+
tracer_pairs = (core_tracer_pairs..., tc_tracer_pairs..., tc_moisture_pairs...)
273251
return (dynamics_pairs..., tracer_pairs...)
274252
end
275253

276254
# This should prep variables that we will dss in
277255
# dss_hyperdiffusion_tendency_pairs
278256
NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
279-
(; hyperdiff, turbconv_model) = p.atmos
257+
(; hyperdiff, turbconv_model, moisture_model, microphysics_model) = p.atmos
280258
isnothing(hyperdiff) && return nothing
281259

282260
(; ᶜ∇²specific_tracers) = p.hyperdiff
@@ -294,8 +272,8 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
294272
# Note: It is more correct to have ρa inside and outside the divergence
295273
@. ᶜ∇²q_totʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_tot))
296274
end
297-
if p.atmos.moisture_model isa NonEquilMoistModel &&
298-
p.atmos.microphysics_model isa Microphysics1Moment
275+
if moisture_model isa NonEquilMoistModel &&
276+
microphysics_model isa Microphysics1Moment
299277
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs) = p.hyperdiff
300278
for j in 1:n
301279
# Note: It is more correct to have ρa inside and outside the divergence
@@ -304,16 +282,10 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
304282
@. ᶜ∇²q_raiʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_rai))
305283
@. ᶜ∇²q_snoʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_sno))
306284
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
285+
elseif moisture_model isa NonEquilMoistModel &&
286+
microphysics_model isa Microphysics2Moment
287+
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs, ᶜ∇²n_liqʲs, ᶜ∇²n_raiʲs) =
288+
p.hyperdiff
317289
for j in 1:n
318290
# Note: It is more correct to have ρa inside and outside the divergence
319291
@. ᶜ∇²q_liqʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_liq))
@@ -331,7 +303,7 @@ end
331303
# This requires dss to have been called on
332304
# variables in dss_hyperdiffusion_tendency_pairs
333305
NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
334-
(; hyperdiff, turbconv_model) = p.atmos
306+
(; hyperdiff, turbconv_model, moisture_model, microphysics_model) = p.atmos
335307
isnothing(hyperdiff) && return nothing
336308

337309
# Rescale the hyperdiffusivity for precipitating species.
@@ -359,41 +331,28 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
359331
(; ᶜ∇²q_totʲs) = p.hyperdiff
360332
for j in 1:n
361333
@. 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)))
334+
ν₄_scalar * wdivₕ(Y.c.sgsʲs.:($$j).ρa * gradₕ(ᶜ∇²q_totʲs.:($$j)))
335+
@. Yₜ.c.sgsʲs.:($$j).q_tot -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_totʲs.:($$j)))
366336
end
367-
if p.atmos.moisture_model isa NonEquilMoistModel &&
368-
p.atmos.microphysics_model isa Microphysics1Moment
337+
if moisture_model isa NonEquilMoistModel &&
338+
microphysics_model isa Microphysics1Moment
369339
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs) = p.hyperdiff
370340
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)))
341+
@. Yₜ.c.sgsʲs.:($$j).q_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
342+
@. Yₜ.c.sgsʲs.:($$j).q_ice -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
375343
@. Yₜ.c.sgsʲs.:($$j).q_rai -=
376344
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j)))
377345
@. Yₜ.c.sgsʲs.:($$j).q_sno -=
378346
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_snoʲs.:($$j)))
379347
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
348+
elseif moisture_model isa NonEquilMoistModel &&
349+
microphysics_model isa Microphysics2Moment
350+
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs, ᶜ∇²n_liqʲs, ᶜ∇²n_raiʲs) =
351+
p.hyperdiff
390352
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)))
353+
@. Yₜ.c.sgsʲs.:($$j).q_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
354+
@. Yₜ.c.sgsʲs.:($$j).q_ice -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
355+
@. Yₜ.c.sgsʲs.:($$j).n_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²n_liqʲs.:($$j)))
397356
@. Yₜ.c.sgsʲs.:($$j).q_rai -=
398357
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j)))
399358
@. Yₜ.c.sgsʲs.:($$j).q_sno -=

0 commit comments

Comments
 (0)