Skip to content

Commit

Permalink
Update to include Frierson et al. (2006) vertical diffusion
Browse files Browse the repository at this point in the history
Refactor to include diffusion tendency in implicit solver
	modified:   src/cache/precomputed_quantities.jl
	modified:   src/cache/temporary_quantities.jl
	modified:   src/prognostic_equations/vertical_diffusion_boundary_layer.jl
	modified:   src/solver/model_getters.jl
	modified:   src/solver/types.jl
	modified:   perf/flame.jl
	modified:   .buildkite/pipeline.yml
	modified:   config/perf_configs/flame_perf_target_frierson.yml
  • Loading branch information
akshaysridhar committed Jan 9, 2024
1 parent a34a468 commit 159ca25
Show file tree
Hide file tree
Showing 8 changed files with 232 additions and 6 deletions.
8 changes: 8 additions & 0 deletions .buildkite/pipeline.yml
Original file line number Diff line number Diff line change
Expand Up @@ -784,6 +784,14 @@ steps:
artifact_paths: "flame_sphere_baroclinic_wave_rhoe_equilmoist_expvdiff/*"
agents:
slurm_mem: 40GB

- label: ":fire: Flame graph: perf target (frierson diffusion)"
command: >
julia --color=yes --project=perf perf/flame.jl
$PERF_CONFIG_PATH/flame_perf_target_frierson.yml
artifact_paths: "flame_perf_target_frierson/*"
agents:
slurm_mem: 48GB

- label: ":fire: Flame graph: perf target (Threaded)"
command: >
Expand Down
14 changes: 14 additions & 0 deletions config/perf_configs/flame_perf_target_frierson.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
vert_diff: "FriersonDiffusion"
max_newton_iters_ode: 2
use_krylov_method: true
use_dynamic_krylov_rtol: false
krylov_rtol: 0.99
precip_model: "0M"
z_elem: 20
dz_bottom: 100
dt_save_state_to_disk: "12hours"
initial_condition: "MoistBaroclinicWave"
dt: "40secs"
t_end: "12hours"
job_id: "flame_perf_target_frierson"
moist: "equil"
11 changes: 6 additions & 5 deletions perf/flame.jl
Original file line number Diff line number Diff line change
Expand Up @@ -37,17 +37,18 @@ ProfileCanvas.html_file(joinpath(output_dir, "flame.html"), results)
#####

allocs_limit = Dict()
allocs_limit["flame_perf_target"] = 178_720
allocs_limit["flame_perf_target_tracers"] = 210_976
allocs_limit["flame_perf_target"] = 190_816
allocs_limit["flame_perf_target_tracers"] = 223_072
allocs_limit["flame_perf_target_edmfx"] = 7_005_552
allocs_limit["flame_perf_diagnostics"] = 26_645_600
allocs_limit["flame_perf_target_diagnostic_edmfx"] = 1_350_976
allocs_limit["flame_perf_target_diagnostic_edmfx"] = 1_380_992
allocs_limit["flame_sphere_baroclinic_wave_rhoe_equilmoist_expvdiff"] =
4_018_252_656
allocs_limit["flame_perf_target_frierson"] = 8_030_478_736
allocs_limit["flame_perf_target_threaded"] = 1_276_864
allocs_limit["flame_perf_target_callbacks"] = 386_584
allocs_limit["flame_perf_gw"] = 3_240_622_560
allocs_limit["flame_perf_target_prognostic_edmfx_aquaplanet"] = 1_274_208
allocs_limit["flame_perf_gw"] = 3_261_869_920
allocs_limit["flame_perf_target_prognostic_edmfx_aquaplanet"] = 1_310_048

# Ideally, we would like to track all the allocations, but this becomes too
# expensive there is too many of them. Here, we set the default sample rate to
Expand Down
190 changes: 190 additions & 0 deletions src/cache/precomputed_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,9 @@ function precomputed_quantities(Y, atmos)
vert_diff_quantities = if atmos.vert_diff isa VerticalDiffusion
ᶜK_h = similar(Y.c, FT)
(; ᶜK_u = ᶜK_h, ᶜK_h) # ᶜK_u aliases ᶜK_h because they are always equal.
elseif atmos.vert_diff isa FriersonDiffusion
ᶜK_h = similar(Y.c, FT)
(; ᶜK_u = ᶜK_h, ᶜK_h) # ᶜK_u aliases ᶜK_h because they are always equal.
else
(;)
end
Expand Down Expand Up @@ -285,6 +288,110 @@ function eddy_diffusivity_coefficient(C_E, norm_v_a, z_a, p)
K_E = C_E * norm_v_a * z_a
return p > p_pbl ? K_E : K_E * exp(-((p_pbl - p) / p_strato)^2)
end
function eddy_diffusivity_coefficient(
z::FT,
z₀,
f_b::FT,
h::FT,
uₐ,
C_E::FT,
Ri::FT,
Ri_a::FT,
Ri_c::FT,
κ::FT,
) where {FT}
# Equations (17), (18)
if z < f_b * h
K_b =
compute_surface_layer_diffusivity(z, z₀, κ, C_E, Ri, Ri_a, Ri_c, uₐ)
return K_b
elseif f_b * h < z < h
K_b = compute_surface_layer_diffusivity(
f_b * h,
z₀,
κ,
C_E,
Ri,
Ri_a,
Ri_c,
uₐ,
)
K = K_b * (z / f_b / h) * (1 - (z - f_b * h) / (1 - f_b) / h)^2
return K
else
return FT(0)
end
end

### Frierson (2006) diffusion
function compute_boundary_layer_height!(
h_boundary_layer,
f_b::FT,
dz,
Ri_local,
Ri_c::FT,
Ri_a,
) where {FT}
Fields.bycolumn(axes(Ri_local)) do colidx
@inbounds for il in 1:Spaces.nlevels(axes(Ri_local[colidx]))
h_boundary_layer[colidx] .=
ifelse.(
Fields.Field(
Fields.field_values(Fields.level(Ri_local[colidx], il)),
axes(h_boundary_layer[colidx]),
) .< Ri_c,
Fields.Field(
Fields.field_values(Fields.level(dz[colidx], il)),
axes(h_boundary_layer[colidx]),
),
h_boundary_layer[colidx],
)
end
end
end

function compute_bulk_richardson_number(
θ_v,
θ_v_a,
norm_ua,
grav,
z::FT,
) where {FT}
return (grav * z) * (θ_v - θ_v_a) / (θ_v_a * (norm_ua)^2 + FT(1))
end
function compute_exchange_coefficient(Ri_a, Ri_c, zₐ, z₀, κ::FT) where {FT}
# Equations (12), (13), (14)
if Ri_a < FT(0)
return κ^2 * (log(zₐ / z₀))^(-2)
elseif FT(0) < Ri_a < Ri_c
return κ^2 * (log(zₐ / z₀))^(-2) * (1 - Ri_a / Ri_c)^2
else
return FT(0)
end
end

function compute_surface_layer_diffusivity(
z::FT,
z₀::FT,
κ::FT,
C_E::FT,
Ri::FT,
Ri_a::FT,
Ri_c::FT,
norm_uₐ,
) where {FT}
# Equations (19), (20)
if Ri_a < FT(0)
return κ * norm_uₐ * sqrt(C_E) * z
else
return κ *
norm_uₐ *
sqrt(C_E) *
z *
(1 + Ri / Ri_c * (log(z / z₀) / (1 - Ri / Ri_c)))^(-1)
end
end
###

"""
set_precomputed_quantities!(Y, p, t)
Expand Down Expand Up @@ -374,6 +481,89 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
ᶜΔz_surface / 2,
ᶜp,
)
elseif vert_diff isa FriersonDiffusion
(; ᶜK_h, sfc_conditions, ᶜts) = p.precomputed
(; params) = p
interior_uₕ = Fields.level(Y.c.uₕ, 1)
κ = CAP.von_karman_const(params)
grav = CAP.grav(params)
FT = Spaces.undertype(axes(ᶜK_h))
z₀ = FT(1e-5)
Ri_c = FT(1.0)
f_b = FT(0.10)

# Prepare scratch vars
ᶠρK_E = p.scratch.ᶠtemp_scalar
θ_v = p.scratch.ᶜtemp_scalar
Ri = p.scratch.ᶜtemp_scalar_2
dz_local = p.scratch.ᶜtemp_scalar_3
θ_v_sfc = p.scratch.ᶠtemp_field_level
Ri_a = p.scratch.temp_field_level
z_local = p.scratch.temp_data
z_sfc = p.scratch.temp_data_face_level
ᶜθ_v_sfc = C_E = p.scratch.temp_field_level_2
h_boundary_layer = p.scratch.temp_field_level_3
ᶠts_sfc = sfc_conditions.ts
ᶜz = Fields.coordinate_field(Y.c).z
interior_uₕ = Fields.level(Y.c.uₕ, 1)
ᶜΔz_surface = Fields.Δz_field(interior_uₕ)
@. θ_v = TD.virtual_pottemp(thermo_params, ᶜts)
@. θ_v_sfc = TD.virtual_pottemp(thermo_params, ᶠts_sfc)
θ_v_a = Fields.level(θ_v, 1)

## Compute boundary layer height

## TODO: Cache elevation field?
z_local .= Fields.field_values(Fields.coordinate_field(Y.c).z)
z_sfc .= Fields.field_values(
Fields.level(Fields.coordinate_field(Y.f).z, half),
)
@. z_local = z_local - z_sfc
dz_local .= Fields.Field(z_local, axes(Y.c))
ᶜθ_v_sfc .=
Fields.Field(Fields.field_values(θ_v_sfc), axes(interior_uₕ))

@. Ri = compute_bulk_richardson_number(
θ_v,
θ_v_a,
norm(Y.c.uₕ),
grav,
dz_local,
)
@. Ri_a = compute_bulk_richardson_number(
θ_v_a,
ᶜθ_v_sfc,
norm(interior_uₕ),
grav,
ᶜΔz_surface / 2,
)

#### Detect 𝒽, boundary layer height per column
h_boundary_layer = f_b .* Fields.level(ᶜz, Spaces.nlevels(axes(Y.c)))
compute_boundary_layer_height!(
h_boundary_layer,
f_b,
dz_local,
Ri,
Ri_c,
Ri_a,
)

## Exchange coefficients
@. C_E =
compute_exchange_coefficient(Ri_a, Ri_c, ᶜΔz_surface ./ 2, z₀, κ)
@. ᶜK_h = eddy_diffusivity_coefficient(
dz_local,
z₀,
f_b,
h_boundary_layer,
norm(interior_uₕ),
C_E,
Ri,
Ri_a,
Ri_c,
κ,
)
end

if precip_model isa Microphysics1Moment
Expand Down
9 changes: 9 additions & 0 deletions src/cache/temporary_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,15 @@ function temporary_quantities(Y, atmos)
ᶠtemp_scalar = Fields.Field(FT, face_space), # ᶠp, ᶠρK_E
ᶜtemp_scalar = Fields.Field(FT, center_space), # ᶜ1
ᶜtemp_scalar_2 = Fields.Field(FT, center_space), # ᶜtke_exch
ᶜtemp_scalar_3 = Fields.Field(FT, center_space),
ᶠtemp_field_level = Fields.level(Fields.Field(FT, face_space), half),
temp_field_level = Fields.level(Fields.Field(FT, center_space), 1),
temp_field_level_2 = Fields.level(Fields.Field(FT, center_space), 1),
temp_field_level_3 = Fields.level(Fields.Field(FT, center_space), 1),
temp_data = Fields.field_values(Fields.Field(FT, center_space)),
temp_data_face_level = Fields.field_values(
Fields.level(Fields.Field(FT, face_space), half),
), # ρaʲu³ʲ_data
temp_data_level = Fields.field_values(
Fields.level(Fields.Field(FT, center_space), 1),
), # ρaʲu³ʲ_data
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ function vertical_diffusion_boundary_layer_tendency!(
p,
t,
colidx,
::VerticalDiffusion,
::Union{VerticalDiffusion, FriersonDiffusion},
)
FT = eltype(Y)
(; ᶜu, ᶜh_tot, ᶜspecific, ᶜK_u, ᶜK_h, sfc_conditions) = p.precomputed
Expand Down
2 changes: 2 additions & 0 deletions src/solver/model_getters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ function get_vertical_diffusion_model(
nothing
elseif vert_diff_name in ("true", true, "VerticalDiffusion")
VerticalDiffusion{diffuse_momentum, FT}(; C_E = params.C_E)
elseif vert_diff_name in ("FriersonDiffusion",)
FriersonDiffusion{diffuse_momentum, FT}()
else
error("Uncaught diffusion model `$vert_diff_name`.")
end
Expand Down
2 changes: 2 additions & 0 deletions src/solver/types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@ Base.@kwdef struct VerticalDiffusion{DM, FT} <: AbstractVerticalDiffusion
C_E::FT
end
diffuse_momentum(::VerticalDiffusion{DM}) where {DM} = DM
struct FriersonDiffusion{DM, FT} <: AbstractVerticalDiffusion end
diffuse_momentum(::FriersonDiffusion{DM}) where {DM} = DM
diffuse_momentum(::Nothing) = false

abstract type AbstractSponge end
Expand Down

0 comments on commit 159ca25

Please sign in to comment.