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
CliMA · Jan 9, 2024 · 159ca25 · 159ca25
1 parent a34a468
commit 159ca25
Show file tree

Hide file tree

Showing 8 changed files with 232 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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

diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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