Skip to content

rft: clean up hyperdiffusion #4049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 24, 2025
Merged

rft: clean up hyperdiffusion #4049

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 48 additions & 89 deletions src/prognostic_equations/hyperdiffusion.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,26 +22,16 @@ function ν₄(hyperdiff, Y)
return (; ν₄_scalar, ν₄_vorticity)
end

hyperdiffusion_cache(Y, atmos) = hyperdiffusion_cache(
Y,
atmos.hyperdiff,
atmos.turbconv_model,
atmos.moisture_model,
atmos.microphysics_model,
)

# No hyperdiffiusion
hyperdiffusion_cache(Y, hyperdiff::Nothing, _, _, _) = (;)
function hyperdiffusion_cache(Y, atmos)
(; hyperdiff, turbconv_model, moisture_model, microphysics_model) = atmos
isnothing(hyperdiff) && return (;) # No hyperdiffiusion
hyperdiffusion_cache(Y, hyperdiff, turbconv_model, moisture_model, microphysics_model)
end

function hyperdiffusion_cache(
Y,
hyperdiff::ClimaHyperdiffusion,
turbconv_model,
moisture_model,
microphysics_model,
Y, hyperdiff::ClimaHyperdiffusion, turbconv_model, moisture_model, microphysics_model,
)
quadrature_style =
Spaces.quadrature_style(Spaces.horizontal_space(axes(Y.c)))
quadrature_style = Spaces.quadrature_style(Spaces.horizontal_space(axes(Y.c)))
FT = eltype(Y)
n = n_mass_flux_subdomains(turbconv_model)

Expand All @@ -54,9 +44,7 @@ function hyperdiffusion_cache(
)

# Sub-grid scale quantities
ᶜ∇²uʲs =
turbconv_model isa PrognosticEDMFX ? similar(Y.c, NTuple{n, C123{FT}}) :
(;)
ᶜ∇²uʲs = turbconv_model isa PrognosticEDMFX ? similar(Y.c, NTuple{n, C123{FT}}) : (;)
moisture_sgs_quantities =
moisture_model isa NonEquilMoistModel &&
microphysics_model isa Microphysics1Moment ?
Expand Down Expand Up @@ -86,17 +74,13 @@ function hyperdiffusion_cache(
moisture_sgs_quantities...,
) : (;)
maybe_ᶜ∇²tke⁰ =
use_prognostic_tke(turbconv_model) ? (; ᶜ∇²tke⁰ = similar(Y.c, FT)) :
(;)
use_prognostic_tke(turbconv_model) ? (; ᶜ∇²tke⁰ = similar(Y.c, FT)) : (;)
sgs_quantities = (; sgs_quantities..., maybe_ᶜ∇²tke⁰...)
quantities = (; gs_quantities..., sgs_quantities...)
if do_dss(axes(Y.c))
quantities = (;
quantities...,
hyperdiffusion_ghost_buffer = map(
Spaces.create_dss_buffer,
quantities,
),
hyperdiffusion_ghost_buffer = map(Spaces.create_dss_buffer, quantities),
)
end
return (; quantities..., ᶜ∇²u, ᶜ∇²uʲs)
Expand All @@ -115,28 +99,25 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)

n = n_mass_flux_subdomains(turbconv_model)
diffuse_tke = use_prognostic_tke(turbconv_model)
(; ᶜp) = p.precomputed
(; ᶜp, ᶜu) = p.precomputed
(; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff
if turbconv_model isa PrognosticEDMFX
(; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
(; ᶜuʲs) = p.precomputed
end

# Grid scale hyperdiffusion
@. ᶜ∇²u =
C123(wgradₕ(divₕ(p.precomputed.ᶜu))) -
C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜu))))
@. ᶜ∇²u = C123(wgradₕ(divₕ(ᶜu))) - C123(wcurlₕ(C123(curlₕ(ᶜu))))

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

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

if diffuse_tke
ᶜρa⁰ =
turbconv_model isa PrognosticEDMFX ?
(@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
ᶜtke⁰ =
@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
(; ᶜ∇²tke⁰) = p.hyperdiff
@. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰))
end
Expand All @@ -145,8 +126,7 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)
if turbconv_model isa PrognosticEDMFX
for j in 1:n
@. ᶜ∇²uʲs.:($$j) =
C123(wgradₕ(divₕ(p.precomputed.ᶜuʲs.:($$j)))) -
C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜuʲs.:($$j)))))
C123(wgradₕ(divₕ(ᶜuʲs.:($$j)))) - C123(wcurlₕ(C123(curlₕ(ᶜuʲs.:($$j)))))
@. ᶜ∇²mseʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).mse))
@. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j))
@. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j))
Expand All @@ -165,11 +145,12 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)

n = n_mass_flux_subdomains(turbconv_model)
diffuse_tke = use_prognostic_tke(turbconv_model)
ᶜρ = Y.c.ρ
ᶜJ = Fields.local_geometry_field(Y.c).J
point_type = eltype(Fields.coordinate_field(Y.c))
(; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff
if turbconv_model isa PrognosticEDMFX
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
ᶜρa⁰ = @. lazy(ρa⁰(ᶜρ, Y.c.sgsʲs, turbconv_model))
(; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
end
if use_prognostic_tke(turbconv_model)
Expand All @@ -183,13 +164,13 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
end

# re-use to store the curl-curl part
@. ᶜ∇²u =
ᶜ∇⁴u = @. ᶜ∇²u =
divergence_damping_factor * C123(wgradₕ(divₕ(ᶜ∇²u))) -
C123(wcurlₕ(C123(curlₕ(ᶜ∇²u))))
@. Yₜ.c.uₕ -= ν₄_vorticity * C12(ᶜ∇²u)
@. Yₜ.f.u₃ -= ν₄_vorticity * ᶠwinterp(ᶜJ * Y.c.ρ, C3(ᶜ∇²u))
@. Yₜ.c.uₕ -= ν₄_vorticity * C12(ᶜ∇⁴u)
@. Yₜ.f.u₃ -= ν₄_vorticity * ᶠwinterp(ᶜJ * ᶜρ, C3(ᶜ∇⁴u))

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

# Sub-grid scale hyperdiffusion continued
if (turbconv_model isa PrognosticEDMFX) && diffuse_tke
Expand All @@ -199,18 +180,16 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
for j in 1:n
if point_type <: Geometry.Abstract3DPoint
# only need curl-curl part
@. ᶜ∇²uᵥʲs.:($$j) = C3(wcurlₕ(C123(curlₕ(ᶜ∇²uʲs.:($$j)))))
@. Yₜ.f.sgsʲs.:($$j).u₃ +=
ν₄_vorticity * ᶠwinterp(ᶜJ * Y.c.ρ, ᶜ∇²uᵥʲs.:($$j))
ᶜ∇⁴uᵥʲ = @. ᶜ∇²uᵥʲs.:($$j) = C3(wcurlₕ(C123(curlₕ(ᶜ∇²uʲs.:($$j)))))
@. Yₜ.f.sgsʲs.:($$j).u₃ += ν₄_vorticity * ᶠwinterp(ᶜJ * ᶜρ, ᶜ∇⁴uᵥʲ)
end
# Note: It is more correct to have ρa inside and outside the divergence
@. Yₜ.c.sgsʲs.:($$j).mse -=
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²mseʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).mse -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²mseʲs.:($$j)))
end
end

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

Expand Down Expand Up @@ -268,15 +247,14 @@ function dss_hyperdiffusion_tendency_pairs(p)
p.hyperdiff.ᶜ∇²q_liqʲs => buffer.ᶜ∇²q_liqʲs,
p.hyperdiff.ᶜ∇²q_raiʲs => buffer.ᶜ∇²q_raiʲs,
) : ()
tracer_pairs =
(core_tracer_pairs..., tc_tracer_pairs..., tc_moisture_pairs...)
tracer_pairs = (core_tracer_pairs..., tc_tracer_pairs..., tc_moisture_pairs...)
return (dynamics_pairs..., tracer_pairs...)
end

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

(; ᶜ∇²specific_tracers) = p.hyperdiff
Expand All @@ -294,8 +272,8 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
# Note: It is more correct to have ρa inside and outside the divergence
@. ᶜ∇²q_totʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_tot))
end
if p.atmos.moisture_model isa NonEquilMoistModel &&
p.atmos.microphysics_model isa Microphysics1Moment
if moisture_model isa NonEquilMoistModel &&
microphysics_model isa Microphysics1Moment
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs) = p.hyperdiff
for j in 1:n
# Note: It is more correct to have ρa inside and outside the divergence
Expand All @@ -304,16 +282,10 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
@. ᶜ∇²q_raiʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_rai))
@. ᶜ∇²q_snoʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_sno))
end
elseif p.atmos.moisture_model isa NonEquilMoistModel &&
p.atmos.microphysics_model isa Microphysics2Moment
(;
ᶜ∇²q_liqʲs,
ᶜ∇²q_iceʲs,
ᶜ∇²q_raiʲs,
ᶜ∇²q_snoʲs,
ᶜ∇²n_liqʲs,
ᶜ∇²n_raiʲs,
) = p.hyperdiff
elseif moisture_model isa NonEquilMoistModel &&
microphysics_model isa Microphysics2Moment
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs, ᶜ∇²n_liqʲs, ᶜ∇²n_raiʲs) =
p.hyperdiff
for j in 1:n
# Note: It is more correct to have ρa inside and outside the divergence
@. ᶜ∇²q_liqʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).q_liq))
Expand All @@ -331,7 +303,7 @@ end
# This requires dss to have been called on
# variables in dss_hyperdiffusion_tendency_pairs
NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
(; hyperdiff, turbconv_model) = p.atmos
(; hyperdiff, turbconv_model, moisture_model, microphysics_model) = p.atmos
isnothing(hyperdiff) && return nothing

# Rescale the hyperdiffusivity for precipitating species.
Expand Down Expand Up @@ -359,41 +331,28 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
(; ᶜ∇²q_totʲs) = p.hyperdiff
for j in 1:n
@. Yₜ.c.sgsʲs.:($$j).ρa -=
ν₄_scalar *
wdivₕ(Y.c.sgsʲs.:($$j).ρa * gradₕ(ᶜ∇²q_totʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_tot -=
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_totʲs.:($$j)))
ν₄_scalar * wdivₕ(Y.c.sgsʲs.:($$j).ρa * gradₕ(ᶜ∇²q_totʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_tot -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_totʲs.:($$j)))
end
if p.atmos.moisture_model isa NonEquilMoistModel &&
p.atmos.microphysics_model isa Microphysics1Moment
if moisture_model isa NonEquilMoistModel &&
microphysics_model isa Microphysics1Moment
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs) = p.hyperdiff
for j in 1:n
@. Yₜ.c.sgsʲs.:($$j).q_liq -=
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_ice -=
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_ice -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_rai -=
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_sno -=
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_snoʲs.:($$j)))
end
elseif p.atmos.moisture_model isa NonEquilMoistModel &&
p.atmos.microphysics_model isa Microphysics2Moment
(;
ᶜ∇²q_liqʲs,
ᶜ∇²q_iceʲs,
ᶜ∇²q_raiʲs,
ᶜ∇²q_snoʲs,
ᶜ∇²n_liqʲs,
ᶜ∇²n_raiʲs,
) = p.hyperdiff
elseif moisture_model isa NonEquilMoistModel &&
microphysics_model isa Microphysics2Moment
(; ᶜ∇²q_liqʲs, ᶜ∇²q_iceʲs, ᶜ∇²q_raiʲs, ᶜ∇²q_snoʲs, ᶜ∇²n_liqʲs, ᶜ∇²n_raiʲs) =
p.hyperdiff
for j in 1:n
@. Yₜ.c.sgsʲs.:($$j).q_liq -=
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_ice -=
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).n_liq -=
ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²n_liqʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_liqʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_ice -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).n_liq -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²n_liqʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_rai -=
ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j)))
@. Yₜ.c.sgsʲs.:($$j).q_sno -=
Expand Down
Loading