Skip to content

Commit

Permalink
Modify velocity, kinetic energy, and curl terms to match dycore paper
Browse files Browse the repository at this point in the history
  • Loading branch information
@dennisYatunin
dennisYatunin committed Jan 30, 2025
1 parent 2612ba7 commit ddb2af6
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 102 deletions.
27 changes: 8 additions & 19 deletions src/cache/diagnostic_edmf_precomputed_quantities.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,21 +9,16 @@ import ClimaCore: Spaces, Fields, RecursiveApply
uₕ_level,
u³_halflevel,
local_geometry_level,
local_geometry_halflevel,
)
g³³_level = g³³(local_geometry_level.gⁱʲ)
uₕʰ_level = CT12(uₕ_level, local_geometry_level)
uₕ³_level = CT3(uₕ_level, local_geometry_level)
u₃³_level = u³_halflevel - uₕ³_level # assume that u³_level = u³_halflevel
u₃_level = inv(g³³_level) * u₃³_level
return (
dot(
C123(uₕ_level, local_geometry_level),
CT123(uₕ_level, local_geometry_level),
) +
dot(
C123(u³_halflevel, local_geometry_halflevel),
CT123(u³_halflevel, local_geometry_halflevel),
) +
2 * dot(
CT123(uₕ_level, local_geometry_level),
C123(u³_halflevel, local_geometry_halflevel),
)
dot(uₕʰ_level, uₕ_level) +
2 * dot(uₕ³_level, u₃_level) +
dot(u₃³_level, u₃_level)
) / 2
end

Expand Down Expand Up @@ -58,7 +53,6 @@ NVTX.@annotate function set_diagnostic_edmfx_env_quantities_level!(
uₕ_level,
K⁰_level,
local_geometry_level,
local_geometry_halflevel,
turbconv_model,
)
@. u³⁰_halflevel = divide_by_ρa(
Expand All @@ -73,7 +67,6 @@ NVTX.@annotate function set_diagnostic_edmfx_env_quantities_level!(
uₕ_level,
u³⁰_halflevel,
local_geometry_level,
local_geometry_halflevel,
)
return nothing
end
Expand Down Expand Up @@ -178,7 +171,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
uₕ_int_level,
u³ʲ_int_halflevel,
local_geometry_int_level,
local_geometry_int_halflevel,
)
set_diagnostic_edmfx_draft_quantities_level!(
thermo_params,
Expand All @@ -205,7 +197,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
uₕ_int_level,
K⁰_int_level,
local_geometry_int_level,
local_geometry_int_halflevel,
turbconv_model,
)

Expand Down Expand Up @@ -798,7 +789,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
uₕ_level,
u³ʲ_halflevel,
local_geometry_level,
local_geometry_halflevel,
)
set_diagnostic_edmfx_draft_quantities_level!(
thermo_params,
Expand All @@ -823,7 +813,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
uₕ_level,
K⁰_level,
local_geometry_level,
local_geometry_halflevel,
turbconv_model,
)
end
Expand Down
83 changes: 42 additions & 41 deletions src/cache/precomputed_quantities.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ function precomputed_quantities(Y, atmos)
ᶜspecific = specific_gs.(Y.c),
ᶜu = similar(Y.c, C123{FT}),
ᶠu³ = similar(Y.f, CT3{FT}),
ᶠu = similar(Y.f, CT123{FT}),
ᶜwₜqₜ = similar(Y.c, Geometry.WVector{FT}),
ᶜwₕhₜ = similar(Y.c, Geometry.WVector{FT}),
ᶜK = similar(Y.c, FT),
Expand Down Expand Up @@ -86,6 +87,7 @@ function precomputed_quantities(Y, atmos)
ᶠu₃⁰ = similar(Y.f, C3{FT}),
ᶜu⁰ = similar(Y.c, C123{FT}),
ᶠu³⁰ = similar(Y.f, CT3{FT}),
ᶠu⁰ = similar(Y.f, CT123{FT}),
ᶜK⁰ = similar(Y.c, FT),
ᶜmse⁰ = similar(Y.c, FT),
ᶜq_tot⁰ = similar(Y.c, FT),
Expand All @@ -97,6 +99,7 @@ function precomputed_quantities(Y, atmos)
ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
ᶜuʲs = similar(Y.c, NTuple{n, C123{FT}}),
ᶠu³ʲs = similar(Y.f, NTuple{n, CT3{FT}}),
ᶠuʲs = similar(Y.f, NTuple{n, CT123{FT}}),
ᶜKʲs = similar(Y.c, NTuple{n, FT}),
ᶠKᵥʲs = similar(Y.f, NTuple{n, FT}),
ᶜtsʲs = similar(Y.c, NTuple{n, TST}),
Expand Down Expand Up @@ -186,29 +189,24 @@ function precomputed_quantities(Y, atmos)
)
end

# Interpolates the third contravariant component of Y.c.uₕ to cell faces.
function set_ᶠuₕ³!(ᶠuₕ³, Y)
ᶜJ = Fields.local_geometry_field(Y.c).J
@. ᶠuₕ³ = ᶠwinterp(Y.c.ρ * ᶜJ, CT3(Y.c.uₕ))
return nothing
end

"""
set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
set_velocity_at_surface!(Y, turbconv_model)
Modifies `Y.f.u₃` so that `ᶠu³` is 0 at the surface. Specifically, since
`u³ = uₕ³ + u³ = uₕ³ + u₃ * g³³`, setting `u³` to 0 gives `u₃ = -uₕ³ / g³³`. If
the `turbconv_model` is EDMFX, the `Y.f.sgsʲs` are also modified so that each
`u₃ʲ` is equal to `u₃` at the surface.
`u³ = g³ʰ uₕ + g³³ u₃`, setting `u³` to 0 gives `u₃ = -(g³³)⁻¹ g³ʰ uₕ`. If the
`turbconv_model` is EDMFX, the `Y.f.sgsʲs` are also modified so that each `u₃ʲ`
is equal to `u₃` at the surface.
"""
function set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
sfc_u₃ = Fields.level(Y.f.u₃.components.data.:1, half)
sfc_uₕ³ = Fields.level(ᶠuₕ³.components.data.:1, half)
sfc_g³³ = g³³_field(sfc_u₃)
@. sfc_u₃ = -sfc_uₕ³ / sfc_g³³ # u³ = uₕ³ + w³ = uₕ³ + w₃ * g³³
function set_velocity_at_surface!(Y, turbconv_model)
bot_uₕ = Fields.level(Y.c.uₕ, 1)
sfc_u₃ = Fields.level(Y.f.u₃, half)
sfc_gⁱʲ = Fields.local_geometry_field(sfc_u₃).gⁱʲ
shifted_sfc_u₃ = Fields.Field(Fields.field_values(sfc_u₃), axes(bot_uₕ))
shifted_sfc_gⁱʲ = Fields.Field(Fields.field_values(sfc_gⁱʲ), axes(bot_uₕ))
@. shifted_sfc_u₃ = -inv(g³³(shifted_sfc_gⁱʲ)) * CT3(bot_uₕ)
if turbconv_model isa PrognosticEDMFX
for j in 1:n_mass_flux_subdomains(turbconv_model)
sfc_u₃ʲ = Fields.level(Y.f.sgsʲs.:($j).u₃.components.data.:1, half)
sfc_u₃ʲ = Fields.level(Y.f.sgsʲs.:($j).u₃, half)
@. sfc_u₃ʲ = sfc_u₃
end
end
Expand All @@ -221,15 +219,12 @@ end
Modifies `Y.f.u₃` so that `u₃` is 0 at the model top.
"""
function set_velocity_at_top!(Y, turbconv_model)
top_u₃ = Fields.level(
Y.f.u₃.components.data.:1,
Spaces.nlevels(axes(Y.c)) + half,
)
@. top_u₃ = 0
top_u₃ = Fields.level(Y.f.u₃, Spaces.nlevels(axes(Y.c)) + half)
@. top_u₃ = C3(0)
if turbconv_model isa PrognosticEDMFX
for j in 1:n_mass_flux_subdomains(turbconv_model)
top_u₃ʲ = Fields.level(
Y.f.sgsʲs.:($j).u₃.components.data.:1,
Y.f.sgsʲs.:($j).u₃,
Spaces.nlevels(axes(Y.c)) + half,
)
@. top_u₃ʲ = top_u₃
Expand All @@ -238,15 +233,21 @@ function set_velocity_at_top!(Y, turbconv_model)
return nothing
end

# This is used to set the grid-scale velocity quantities ᶜu, ᶠu³, ᶜK based on
# ᶠu₃, and it is also used to set the SGS quantities based on ᶠu₃⁰ and ᶠu₃ʲ.
function set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, ᶠu₃, ᶜuₕ, ᶠuₕ³)
@. ᶜu = C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
@. ᶠu³ = ᶠuₕ³ + CT3(ᶠu₃)
function set_ᶜu_and_ᶜK!(ᶜu, ᶜK, ᶜuₕ, ᶠu₃)
@. ᶜu = C123(ᶜuₕ) + C123(ᶜinterp(ᶠu₃))
compute_kinetic!(ᶜK, ᶜuₕ, ᶠu₃)
return nothing
end

function set_ᶠu³_and_ᶠu!(ᶠu³, ᶠu, ᶜu, ᶠu₃, ᶜρ)
ᶜJ = Fields.local_geometry_field(ᶜu).J
ᶜgⁱʲ = Fields.local_geometry_field(ᶜu).gⁱʲ
@. ᶠu³ =
ᶠwinterp(ᶜρ * ᶜJ, CT3(C12(ᶜu))) + ᶠwinterp(ᶜρ * ᶜJ, g³³(ᶜgⁱʲ)) * ᶠu₃
@. ᶠu = CT123(ᶠwinterp(ᶜρ * ᶜJ, CT12(ᶜu))) + CT123(ᶠu³)
return nothing
end

function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
ρaʲs(sgsʲs) = map(sgsʲ -> sgsʲ.ρa, sgsʲs)
u₃ʲs(sgsʲs) = map(sgsʲ -> sgsʲ.u₃, sgsʲs)
Expand Down Expand Up @@ -472,18 +473,17 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
n = n_mass_flux_subdomains(turbconv_model)
thermo_args = (thermo_params, moisture_model)
(; ᶜΦ) = p.core
(; ᶜspecific, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp) = p.precomputed
ᶠuₕ³ = p.scratch.ᶠtemp_CT3
(; ᶜspecific, ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp) = p.precomputed

@. ᶜspecific = specific_gs(Y.c)
set_ᶠuₕ³!(ᶠuₕ³, Y)

# TODO: We might want to move this to dss! (and rename dss! to something
# like enforce_constraints!).
set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
set_velocity_at_surface!(Y, turbconv_model)
set_velocity_at_top!(Y, turbconv_model)

set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, Y.f.u₃, Y.c.uₕ, ᶠuₕ³)
set_ᶜu_and_ᶜK!(ᶜu, ᶜK, Y.c.uₕ, Y.f.u₃)
set_ᶠu³_and_ᶠu!(ᶠu³, ᶠu, ᶜu, Y.f.u₃, Y.c.ρ)
if n > 0
# TODO: In the following increments to ᶜK, we actually need to add
# quantities of the form ᶜρaχ⁰ / ᶜρ⁰ and ᶜρaχʲ / ᶜρʲ to ᶜK, rather than
Expand Down Expand Up @@ -550,8 +550,8 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
end

if turbconv_model isa PrognosticEDMFX
set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, t)
set_prognostic_edmf_precomputed_quantities_environment!(Y, p, t)
set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
set_prognostic_edmf_precomputed_quantities_precipitation!(
Y,
Expand Down Expand Up @@ -703,21 +703,22 @@ function output_prognostic_sgs_quantities(Y, p, t)
(; turbconv_model) = p.atmos
thermo_params = CAP.thermodynamics_params(p.params)
(; ᶜρa⁰, ᶜρ⁰, ᶜtsʲs) = p.precomputed
ᶠuₕ³ = p.scratch.ᶠtemp_CT3
set_ᶠuₕ³!(ᶠuₕ³, Y)
(ᶠu₃⁺, ᶜu⁺, ᶠu³⁺, ᶜK⁺) =
(ᶠu₃⁺, ᶜu⁺, ᶠu³⁺, ᶠu⁺, ᶜK⁺) =
similar.((
p.precomputed.ᶠu₃⁰,
p.precomputed.ᶜu⁰,
p.precomputed.ᶠu³⁰,
p.precomputed.ᶠu⁰,
p.precomputed.ᶜK⁰,
))
set_sgs_ᶠu₃!(u₃⁺, ᶠu₃⁺, Y, turbconv_model)
set_velocity_quantities!(ᶜu⁺, ᶠu³⁺, ᶜK⁺, ᶠu₃⁺, Y.c.uₕ, ᶠuₕ³)
set_ᶜu_and_ᶜK!(ᶜu⁺, ᶜK⁺, Y.c.uₕ, ᶠu₃⁺)
ᶜts⁺ = ᶜtsʲs.:1
ᶜa⁺ = @. draft_area(ρa⁺(Y.c), TD.air_density(thermo_params, ᶜts⁺))
ᶜρ⁺ = TD.air_density.(thermo_params, ᶜts⁺)
set_ᶠu³_and_ᶠu!(ᶠu³⁺, ᶠu⁺, ᶜu⁺, ᶠu₃⁺, ᶜρ⁺)
ᶜa⁺ = @. draft_area(ρa⁺(Y.c), ᶜρ⁺)
ᶜa⁰ = @. draft_area(ᶜρa⁰, ᶜρ⁰)
return (; ᶠu₃⁺, ᶜu⁺, ᶠu³⁺, ᶜK⁺, ᶜts⁺, ᶜa⁺, ᶜa⁰)
return (; ᶠu₃⁺, ᶜu⁺, ᶠu³⁺, ᶠu⁺, ᶜK⁺, ᶜts⁺, ᶜa⁺, ᶜa⁰)
end

"""
Expand Down
21 changes: 10 additions & 11 deletions src/cache/prognostic_edmf_precomputed_quantities.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,13 @@ import Thermodynamics as TD
import ClimaCore: Spaces, Fields

"""
set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
set_prognostic_edmf_precomputed_quantities!(Y, p, t)
Updates the edmf environment precomputed quantities stored in `p` for edmfx.
"""
NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
Y,
p,
ᶠuₕ³,
t,
)
@assert !(p.atmos.moisture_model isa DryModel)
Expand All @@ -22,7 +21,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
(; turbconv_model) = p.atmos
(; ᶜΦ,) = p.core
(; ᶜp, ᶜh_tot, ᶜK) = p.precomputed
(; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
(; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶠu⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
p.precomputed

@. ᶜρa⁰ = ρa⁰(Y.c)
Expand All @@ -42,23 +41,23 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
turbconv_model,
)
set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
set_ᶜu_and_ᶜK!(ᶜu⁰, ᶜK⁰, Y.c.uₕ, ᶠu₃⁰)
# @. ᶜK⁰ += ᶜtke⁰
@. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
@. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
set_ᶠu³_and_ᶠu!(ᶠu³⁰, ᶠu⁰, ᶜu⁰, ᶠu₃⁰, ᶜρ⁰)
return nothing
end

"""
set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, t)
Updates the draft thermo state and boundary conditions
precomputed quantities stored in `p` for edmfx.
"""
NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(
Y,
p,
ᶠuₕ³,
t,
)
(; moisture_model, turbconv_model) = p.atmos
Expand All @@ -74,12 +73,13 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!

(; ᶜΦ,) = p.core
(; ᶜspecific, ᶜp, ᶜh_tot, ᶜK) = p.precomputed
(; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
(; ᶜuʲs, ᶠu³ʲs, ᶠuʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
(; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions

for j in 1:n
ᶜuʲ = ᶜuʲs.:($j)
ᶠu³ʲ = ᶠu³ʲs.:($j)
ᶠuʲ = ᶠuʲs.:($j)
ᶜKʲ = ᶜKʲs.:($j)
ᶠKᵥʲ = ᶠKᵥʲs.:($j)
ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃
Expand All @@ -88,10 +88,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
ᶜmseʲ = Y.c.sgsʲs.:($j).mse
ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot

set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³)
set_ᶜu_and_ᶜK!(ᶜuʲ, ᶜKʲ, Y.c.uₕ, ᶠu₃ʲ)
@. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
@. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
@. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
set_ᶠu³_and_ᶠu!(ᶠu³ʲ, ᶠuʲ, ᶜuʲ, ᶠu₃ʲ, ᶜρʲ)

# EDMFX boundary condition:

Expand Down Expand Up @@ -195,7 +196,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
FT = eltype(params)
n = n_mass_flux_subdomains(turbconv_model)

(; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
(; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶠu⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
(;
ᶜmixing_length_tuple,
ᶜmixing_length,
Expand Down Expand Up @@ -303,8 +304,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
)

# TODO: Currently the shear production only includes vertical gradients
ᶠu⁰ = p.scratch.ᶠtemp_C123
@. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
compute_strain_rate_center!(ᶜstrain_rate, ᶠu⁰)
@. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
Expand Down
Loading

0 comments on commit ddb2af6

Please sign in to comment.