Skip to content
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

Improve design of nonorographic gravity wave parameterization #3191

Merged
merged 2 commits into from
Jul 19, 2024
Merged

Improve design of nonorographic gravity wave parameterization #3191

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d688218
refactor gravity wave cache and tendency
Jul 12, 2024
e2b3287
change function names
Jul 17, 2024
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
156 changes: 109 additions & 47 deletions src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import ClimaCore.Spaces as Spaces
import ClimaCore.Fields as Fields
import ClimaCore.Geometry as Geometry
import ClimaCore.Operators as Operators

non_orographic_gravity_wave_cache(Y, atmos::AtmosModel) =
non_orographic_gravity_wave_cache(
Expand All @@ -27,12 +28,18 @@ function non_orographic_gravity_wave_cache(

nc = Int(floor(FT(2 * cmax / dc + 1)))
c = [FT((n - 1) * dc - cmax) for n in 1:nc]

source_level_z = similar(Fields.level(Y.c.ρ, 1), Tuple{FT, FT})
ᶜlevel = similar(Y.c.ρ, FT)
for i in 1:Spaces.nlevels(axes(Y.c.ρ))
fill!(Fields.level(ᶜlevel, i), i)
end
damp_level_z = similar(source_level_z)
return (;
gw_source_height = source_height,
gw_source_ampl = Bt_0 .* ones(FT, axes(Fields.level(Y.c.ρ, 1))),
gw_Bw = Bw .* ones(FT, axes(Fields.level(Y.c.ρ, 1))),
gw_Bn = Bn .* ones(FT, axes(Fields.level(Y.c.ρ, 1))),
gw_B0 = similar(c),
gw_c = c,
gw_cw = cw .* ones(FT, axes(Fields.level(Y.c.ρ, 1))),
gw_cn = cn .* ones(FT, axes(Fields.level(Y.c.ρ, 1))),
Expand All @@ -41,6 +48,15 @@ function non_orographic_gravity_wave_cache(
gw_nk = Int(nk),
ᶜbuoyancy_frequency = similar(Y.c.ρ),
ᶜdTdz = similar(Y.c.ρ),
source_level_z,
damp_level_z,
source_level = similar(Fields.level(Y.c.ρ, 1)),
damp_level = similar(Fields.level(Y.c.ρ, 1)),
ᶜlevel,
u_phy = similar(Y.c.ρ),
v_phy = similar(Y.c.ρ),
uforcing = similar(Y.c.ρ),
vforcing = similar(Y.c.ρ),
)
end

Expand All @@ -66,6 +82,13 @@ function non_orographic_gravity_wave_cache(
gw_Bw = ones(FT, axes(lat)) .* Bw
gw_cn = ones(FT, axes(lat)) .* cn

source_level_z = similar(Fields.level(Y.c.ρ, 1), Tuple{FT, FT})
ᶜlevel = similar(Y.c.ρ, FT)
for i in 1:Spaces.nlevels(axes(Y.c.ρ))
fill!(Fields.level(ᶜlevel, i), i)
end
damp_level_z = similar(source_level_z)

# This is GFDL source specs -> a smooth function
# source_ampl = @. Bt_0 +
# Bt_n * FT(0.5) * (FT(1) + tanh((lat - ϕ0_n) / dϕ_n)) +
Expand Down Expand Up @@ -94,6 +117,7 @@ function non_orographic_gravity_wave_cache(
gw_source_ampl = source_ampl,
gw_Bw = gw_Bw,
gw_Bn = gw_Bn,
gw_B0 = similar(c),
gw_c = c,
gw_cw = gw_cw,
gw_cn = gw_cn,
Expand All @@ -102,6 +126,15 @@ function non_orographic_gravity_wave_cache(
gw_nk = Int(nk),
ᶜbuoyancy_frequency = similar(Y.c.ρ),
ᶜdTdz = similar(Y.c.ρ),
source_level_z,
damp_level_z,
source_level = similar(Fields.level(Y.c.ρ, 1)),
damp_level = similar(Fields.level(Y.c.ρ, 1)),
ᶜlevel,
u_phy = similar(Y.c.ρ),
v_phy = similar(Y.c.ρ),
uforcing = similar(Y.c.ρ),
vforcing = similar(Y.c.ρ),
)
end

Expand All @@ -116,12 +149,25 @@ function non_orographic_gravity_wave_tendency!(
(; ᶜT,) = p.core
(; ᶜts) = p.precomputed
(; params) = p
(; ᶜdTdz, ᶜbuoyancy_frequency) = p.non_orographic_gravity_wave
(;
ᶜdTdz,
ᶜbuoyancy_frequency,
source_level,
source_level_z,
damp_level,
damp_level_z,
u_phy,
v_phy,
uforcing,
vforcing,
ᶜlevel,
) = p.non_orographic_gravity_wave
(; model_config) = p.atmos
(;
gw_source_ampl,
gw_Bw,
gw_Bn,
gw_B0,
gw_c,
gw_cw,
gw_cn,
Expand All @@ -145,10 +191,10 @@ function non_orographic_gravity_wave_tendency!(
# compute buoyancy frequency
@. ᶜT = TD.air_temperature(thermo_params, ᶜts)

parent(ᶜdTdz) .= parent(Geometry.WVector.(ᶜgradᵥ.(ᶠinterp.(ᶜT))))
ᶜdTdz .= Geometry.WVector.(ᶜgradᵥ.(ᶠinterp.(ᶜT))).components.data.:1
szy21 marked this conversation as resolved.
Show resolved Hide resolved

ᶜbuoyancy_frequency =
@. (grav / ᶜT) * (ᶜdTdz + grav / TD.cp_m(thermo_params, ᶜts))
@. ᶜbuoyancy_frequency =
(grav / ᶜT) * (ᶜdTdz + grav / TD.cp_m(thermo_params, ᶜts))
ᶜbuoyancy_frequency = @. ifelse(
ᶜbuoyancy_frequency < FT(2.5e-5),
FT(sqrt(2.5e-5)),
Expand All @@ -157,56 +203,65 @@ function non_orographic_gravity_wave_tendency!(

if model_config isa SingleColumnModel
# source level: the index of the level that is closest to the source height
source_level = similar(Fields.level(Y.c.ρ, 1))
Fields.bycolumn(axes(ᶜρ)) do colidx
parent(source_level[colidx]) .=
argmin(abs.(parent(ᶜz[colidx]) .- gw_source_height))[1]
end
# damp level: for now we only deposit to top level for column setup
damp_level = similar(Fields.level(Y.c.ρ, 1))
Fields.bycolumn(axes(ᶜρ)) do colidx
parent(damp_level[colidx]) .= length(parent(ᶜz[colidx]))

Operators.column_mapreduce!(
reduce_fun1,
source_level_z,
ᶜz,
ᶜlevel,
) do z, level
(abs.(z .- gw_source_height), level)
end
source_level = source_level_z.:2

Operators.column_mapreduce!(sign, +, damp_level, ᶜz)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dennisYatunin Would something like @. damp_level = Spaces.nlevels(axes(ᶜz)) work?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried it and found that @. damp_level = Spaces.nlevels(axes(ᶜz)) could not work, but we can use fill!(damp_level, Spaces.nlevels(axes(ᶜz))) instead.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, damp_level .= Spaces.nlevels(axes(ᶜz)) or fill!(damp_level, Spaces.nlevels(axes(ᶜz)) would be better here


elseif model_config isa SphericalModel
(; ᶜp) = p.precomputed
# source level: the index of the highest level whose pressure is higher than source pressure
source_level = similar(Fields.level(Y.c.ρ, 1))
Fields.bycolumn(axes(ᶜρ)) do colidx
parent(source_level[colidx]) .=
findlast(parent(ᶜp[colidx]) .> gw_source_pressure)[1]

Operators.column_mapreduce!(
reduce_fun2,
source_level_z,
ᶜp,
ᶜlevel,
) do p, level
(p .- gw_source_pressure, level)
end
source_level = source_level_z.:2


# damp level: the index of the lowest level whose pressure is lower than the damp pressure
damp_level = similar(Fields.level(Y.c.ρ, 1))
Fields.bycolumn(axes(ᶜρ)) do colidx
if sum(parent(ᶜp[colidx]) .< gw_damp_pressure) == 0
parent(damp_level[colidx]) .= length(parent(ᶜz[colidx]))
else
parent(damp_level[colidx]) .=
findfirst(parent(ᶜp[colidx]) .< gw_damp_pressure)[1]
end

Operators.column_mapreduce!(
reduce_fun3,
damp_level_z,
ᶜp,
ᶜlevel,
) do p, level
(p .- gw_damp_pressure, level)
end
damp_level = damp_level_z.:2

end

# prepare physical uv input variables for gravity_wave_forcing()
u_phy = Geometry.UVVector.(Y.c.uₕ).components.data.:1
v_phy = Geometry.UVVector.(Y.c.uₕ).components.data.:2

# a place holder to store physical forcing on uv
uforcing = ones(axes(u_phy))
vforcing = ones(axes(u_phy))

# GW parameterization applied bycolume
Fields.bycolumn(axes(ᶜρ)) do colidx
parent(uforcing[colidx]) .= non_orographic_gravity_wave_forcing(
copy(vec(parent(u_phy[colidx]))),
copy(vec(parent(ᶜbuoyancy_frequency[colidx]))),
copy(vec(parent(ᶜρ[colidx]))),
copy(vec(parent(ᶜz[colidx]))),
vec(parent(u_phy[colidx])),
vec(parent(ᶜbuoyancy_frequency[colidx])),
vec(parent(ᶜρ[colidx])),
vec(parent(ᶜz[colidx])),
Int(parent(source_level[colidx])[1]),
Int(parent(damp_level[colidx])[1]),
parent(gw_source_ampl[colidx])[1],
parent(gw_Bw[colidx])[1],
parent(gw_Bn[colidx])[1],
gw_B0,
parent(gw_cw[colidx])[1],
parent(gw_cn[colidx])[1],
parent(gw_flag[colidx])[1],
Expand All @@ -216,15 +271,16 @@ function non_orographic_gravity_wave_tendency!(
)

parent(vforcing[colidx]) .= non_orographic_gravity_wave_forcing(
copy(vec(parent(v_phy[colidx]))),
copy(vec(parent(ᶜbuoyancy_frequency[colidx]))),
copy(vec(parent(ᶜρ[colidx]))),
copy(vec(parent(ᶜz[colidx]))),
vec(parent(v_phy[colidx])),
vec(parent(ᶜbuoyancy_frequency[colidx])),
vec(parent(ᶜρ[colidx])),
vec(parent(ᶜz[colidx])),
Int(parent(source_level[colidx])[1]),
Int(parent(damp_level[colidx])[1]),
parent(gw_source_ampl[colidx])[1],
parent(gw_Bw)[1],
parent(gw_Bn)[1],
gw_B0,
parent(gw_cw)[1],
parent(gw_cn)[1],
parent(gw_flag)[1],
Expand All @@ -242,33 +298,35 @@ function non_orographic_gravity_wave_tendency!(
end

function non_orographic_gravity_wave_forcing(
ᶜu,
ᶜbf,
ᶜρ,
ᶜz,
old_ᶜu,
old_ᶜbf,
old_ᶜρ,
old_ᶜz,
source_level,
damp_level,
source_ampl,
Bw,
Bn,
B0,
cw,
cn,
flag,
c,
c0,
nk,
)
FT = eltype(ᶜz)
FT = eltype(old_ᶜz)
# add an extra layer above model top so that forcing between the very top
# model layer and the upper boundary can be calculated
append!(ᶜu, FT(2) * ᶜu[end] - ᶜu[end - 1])
append!(ᶜρ, ᶜρ[end] * ᶜρ[end] / ᶜρ[end - 1])
append!(ᶜbf, ᶜbf[end])
append!(ᶜz, FT(2) * ᶜz[end] - ᶜz[end - 1])
ᶜu = vcat(old_ᶜu, FT(2) * old_ᶜu[end] - old_ᶜu[end - 1])
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to note, vcat is not really fixing the issue here, but the change is fine.

ᶜρ = vcat(old_ᶜρ, old_ᶜρ[end] * old_ᶜρ[end] / old_ᶜρ[end - 1])
ᶜbf = vcat(old_ᶜbf, old_ᶜbf[end])
ᶜz = vcat(old_ᶜz, FT(2) * old_ᶜz[end] - old_ᶜz[end - 1])

# wave spectra and the source amplitude
nc = length(c)
c_hat0 = c .- ᶜu[source_level] # c0mu0

Bw_exp = @. exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cw)^2)
Bn_exp = @. exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cn)^2)
B0 = @. sign(c_hat0) * (Bw * Bw_exp + Bn * Bn_exp)
Expand Down Expand Up @@ -388,3 +446,7 @@ end
function calc_intermitency(ρ_source_level, source_ampl, nk, Bsum)
return (source_ampl / ρ_source_level / nk) / Bsum
end

@inline reduce_fun1(a, b) = ifelse(a[1] < b[1], a, b)
@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)
@inline reduce_fun2(a, b) = ifelse(b[1] <= 0, a, b)

(I think this is exactly the same as the previous logic, although it probably doesn't matter much)

@inline reduce_fun3(a, b) = ifelse(a[1] > 0, b, a)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you come up with some more informative names for these functions?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is these names OK?:
@inline min_distance(a, b) = ifelse(a[1] < b[1], a, b)
@inline positive_selector(a, b) = ifelse(b[1] < 0, a, b)
@inline negative_selector(a, b) = ifelse(a[1] > 0, b, a)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
@inline reduce_fun3(a, b) = ifelse(a[1] > 0, b, a)
@inline reduce_fun3(a, b) = ifelse(a[1] >= 0, b, a)

3 changes: 3 additions & 0 deletions test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_3d.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,8 @@ center_u_zonalave = mean(center_u, dims = 1)[1, :, :, :]
center_bf_zonalave = mean(center_bf, dims = 1)[1, :, :, :]
center_ρ_zonalave = mean(center_ρ, dims = 1)[1, :, :, :]

B0 = similar(params.gw_c)

# Jan
month = Dates.month.(time)

Expand All @@ -140,6 +142,7 @@ for j in 1:length(lat)
params.gw_source_ampl,
params.gw_Bw,
params.gw_Bn,
B0,
params.gw_cw,
params.gw_cn,
params.gw_flag,
Expand Down
4 changes: 3 additions & 1 deletion test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_mima.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ bf = @. ifelse(bf < 2.5e-5, sqrt(2.5e-5), sqrt(abs(bf)))

# compute u/v forcings from convective gravity waves
params = non_orographic_gravity_wave(lat, FT)

B0 = similar(params.gw_c)
# nogw forcing

kmax = length(pfull) - 1
Expand Down Expand Up @@ -156,6 +156,7 @@ for i in 1:length(lon)
params.gw_source_ampl[j],
params.gw_Bw,
params.gw_Bn[j],
B0,
params.gw_cw[j],
params.gw_cn,
params.gw_flag[j],
Expand Down Expand Up @@ -190,6 +191,7 @@ for i in 1:length(lon)
params.gw_source_ampl[j],
params.gw_Bw,
params.gw_Bn[j],
B0,
params.gw_cw[j],
params.gw_cn,
params.gw_flag[j],
Expand Down
Loading
Loading