From 81f5fa511f45dd4ed26ae390c46e852c9a6b52d5 Mon Sep 17 00:00:00 2001 From: akshaysridhar Date: Mon, 8 Jan 2024 15:54:35 -0800 Subject: [PATCH] Update to include Frierson et al. (2006) vertical diffusion 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 --- .buildkite/pipeline.yml | 8 + .../flame_perf_target_frierson.yml | 14 ++ perf/flame.jl | 11 +- src/cache/precomputed_quantities.jl | 190 ++++++++++++++++++ src/cache/temporary_quantities.jl | 9 + .../vertical_diffusion_boundary_layer.jl | 2 +- src/solver/model_getters.jl | 2 + src/solver/types.jl | 2 + 8 files changed, 232 insertions(+), 6 deletions(-) create mode 100644 config/perf_configs/flame_perf_target_frierson.yml diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index cf74fb647de..2b5d5194796 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -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: > diff --git a/config/perf_configs/flame_perf_target_frierson.yml b/config/perf_configs/flame_perf_target_frierson.yml new file mode 100644 index 00000000000..f878585b75d --- /dev/null +++ b/config/perf_configs/flame_perf_target_frierson.yml @@ -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" diff --git a/perf/flame.jl b/perf/flame.jl index fc74ad1e917..a26b3362287 100644 --- a/perf/flame.jl +++ b/perf/flame.jl @@ -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_839_008 +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 diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index a61da80f166..2a20d872fd3 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -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 @@ -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) @@ -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 diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl index e42de12ba6b..a0dd921e152 100644 --- a/src/cache/temporary_quantities.jl +++ b/src/cache/temporary_quantities.jl @@ -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 diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index 985740ab3fc..3478bf9838e 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -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 diff --git a/src/solver/model_getters.jl b/src/solver/model_getters.jl index bb4ab881599..fd5b2e4ef98 100644 --- a/src/solver/model_getters.jl +++ b/src/solver/model_getters.jl @@ -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 diff --git a/src/solver/types.jl b/src/solver/types.jl index 8fc5eaf7477..b9773d8d471 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -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