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
+ Formatter

Where allocations are exceeded by less than 6%, increases the limit.
	modified:   perf/flame.jl

perf/flame.jl

@@ -37,17 +37,17 @@ ProfileCanvas.html_file(joinpath(output_dir, "flame.html"), results)
 #####
 
 allocs_limit = Dict()
-allocs_limit["flame_perf_target"] = 148_256
-allocs_limit["flame_perf_target_tracers"] = 180_512
+allocs_limit["flame_perf_target"] = 157_088
+allocs_limit["flame_perf_target_tracers"] = 189_344
 allocs_limit["flame_perf_target_edmfx"] = 7_005_552
 allocs_limit["flame_perf_diagnostics"] = 25_356_928
-allocs_limit["flame_perf_target_diagnostic_edmfx"] = 1_311_040
+allocs_limit["flame_perf_target_diagnostic_edmfx"] = 1_335_232
 allocs_limit["flame_sphere_baroclinic_wave_rhoe_equilmoist_expvdiff"] =
     4_018_252_656
 allocs_limit["flame_perf_target_threaded"] = 1_276_864
 allocs_limit["flame_perf_target_callbacks"] = 37_277_112
-allocs_limit["flame_perf_gw"] = 3_226_429_472
-allocs_limit["flame_perf_target_prognostic_edmfx_aquaplanet"] = 1_258_848
+allocs_limit["flame_perf_gw"] = 3_247_675_296
+allocs_limit["flame_perf_target_prognostic_edmfx_aquaplanet"] = 1_289_568
 
 # 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

src/cache/precomputed_quantities.jl

@@ -113,6 +113,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
@@ -287,6 +290,111 @@ 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,
+    dz,
+    Ri_local,
+    Ri_c::FT,
+) 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]),
+                    ),
+                    Fields.Field(
+                        Fields.field_values(Fields.level(dz[colidx], 1)),
+                        axes(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 + eps(FT))
+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)
@@ -370,6 +478,81 @@ 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(1000.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
+        compute_boundary_layer_height!(h_boundary_layer, dz_local, Ri, Ri_c)
+
+        ## 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

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

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

src/solver/model_getters.jl

@@ -63,6 +63,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

src/solver/types.jl

@@ -37,6 +37,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