updates

Author: haakon-e

config/default_configs/default_config.yml

@@ -163,7 +163,7 @@ vert_diff:
 hyperdiff:
   help: "Hyperdiffusion. Use `CAM_SE` for sensible default values and ClimaHyperdiffusion for user control. [`CAM_SE` (default), `ClimaHyperdiffusion` (or `true`), `none` (or `false`)]"
   value: "CAM_SE"
-smagorinsky:
+smagorinsky_lilly:
   help: "Smagorinsky diffusive closure in the horizontal and vertical [`false` (default), `true`]"
   value: false
 smagorinsky_horizontal:

src/cache/precomputed_quantities.jl

@@ -245,15 +245,10 @@ function precomputed_quantities(Y, atmos)
             precipitation_sgs_quantities...,
             diagnostic_precipitation_sgs_quantities...,
         ) : (;)
-    smagorinsky_lilly_quantities =
-        if atmos.smagorinsky_lilly isa SmagorinskyLilly
-            uvw_vec = UVW(FT(0), FT(0), FT(0))
-            (;
-                ᶜτ_smag = similar(Y.c, typeof(uvw_vec * uvw_vec')),
-                ᶠτ_smag = similar(Y.f, typeof(uvw_vec * uvw_vec')),
-                ᶜD_smag = similar(Y.c, FT),
-                ᶠD_smag = similar(Y.f, FT),
-            )
+    smagorinsky_quantities =
+        if atmos.smagorinsky_horizontal isa SmagorinskyLilly ||
+           atmos.smagorinsky_vertical isa SmagorinskyLilly
+            (; ᶜL_h = similar(Y.c, FT), ᶜL_v = similar(Y.c, FT))
         else
             (;)
         end
@@ -269,7 +264,7 @@ function precomputed_quantities(Y, atmos)
         precipitation_quantities...,
         surface_precip_fluxes...,
         cloud_diagnostics_tuple,
-        smagorinsky_lilly_quantities...,
+        smagorinsky_quantities...,
     )
 end
 
@@ -596,7 +591,8 @@ NVTX.@annotate function set_explicit_precomputed_quantities_part2!(Y, p, t)
         set_cloud_fraction!(Y, p, moisture_model, cloud_model)
     end
 
-    if p.atmos.smagorinsky_lilly isa SmagorinskyLilly
+    if p.atmos.smagorinsky_horizontal isa SmagorinskyLilly ||
+       p.atmos.smagorinsky_vertical isa SmagorinskyLilly
         set_smagorinsky_lilly_precomputed_quantities!(Y, p)
     end
 

src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl

@@ -9,7 +9,7 @@ import ClimaCore: Geometry
 """
     set_smagorinsky_lilly_precomputed_quantities!(Y, p)
 
-Compute the Smagorinsky-Lilly diffusivity tensors, `ᶜτ_smag`, `ᶠτ_smag`, `ᶜD_smag`, and `ᶠD_smag`. 
+Compute the Smagorinsky-Lilly horizontal and vertical length scales `ᶜL_h` and `ᶜL_v`.
 Store in the precomputed quantities `p.precomputed`.
 
 The subgrid-scale momentum flux tensor is defined by `τ = -2 νₜ ∘ S`, where `νₜ` is the Smagorinsky-Lilly eddy viscosity 
@@ -26,66 +26,33 @@ These quantities are computed for both cell centers and faces, with prefixes `
 - `p`: The model parameters, e.g. `AtmosCache`.
 """
 function set_smagorinsky_lilly_precomputed_quantities!(Y, p)
-
-    (; atmos, precomputed, scratch, params) = p
-    c_smag = CAP.c_smag(params)
-    Pr_t = CAP.Prandtl_number_0(CAP.turbconv_params(params))
-    (; ᶜu, ᶠu³, ᶜts, ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶠD_smag) = precomputed
     FT = eltype(Y)
-    grav = CAP.grav(params)
-    thermo_params = CAP.thermodynamics_params(params)
-    (; ᶜtemp_UVWxUVW, ᶠtemp_UVWxUVW, ᶜtemp_strain, ᶠtemp_strain) = scratch
-    (; ᶜtemp_scalar, ᶜtemp_scalar_2, ᶠtemp_scalar, ᶜtemp_UVW, ᶠtemp_UVW) = scratch
-
-    ∇ᵥuvw_boundary = Geometry.outer(Geometry.WVector(0), UVW(0, 0, 0))
-    ᶠgradᵥ_uvw = Operators.GradientC2F(
-        bottom = Operators.SetGradient(∇ᵥuvw_boundary),
-        top = Operators.SetGradient(∇ᵥuvw_boundary),
-    )
-    axis_uvw = (Geometry.UVWAxis(),)
-
-    # Compute UVW velocities
-    ᶜu_uvw = @. ᶜtemp_UVW = UVW(ᶜu)
-    ᶠu_uvw = @. ᶠtemp_UVW = UVW(ᶠinterp(Y.c.uₕ)) + UVW(ᶠu³)
-
-    # Gradients
-    ## cell centers
-    ∇ᶜu_uvw = @. ᶜtemp_UVWxUVW = Geometry.project(axis_uvw, ᶜgradᵥ(ᶠu_uvw))  # vertical component
-    @. ∇ᶜu_uvw += Geometry.project(axis_uvw, gradₕ(ᶜu_uvw))  # horizontal component
-    ## cell faces
-    ∇ᶠu_uvw = @. ᶠtemp_UVWxUVW = Geometry.project(axis_uvw, ᶠgradᵥ_uvw(ᶜu_uvw))  # vertical component
-    @. ∇ᶠu_uvw += Geometry.project(axis_uvw, gradₕ(ᶠu_uvw))  # horizontal component
-
+    (; ᶠu, ᶜts, ᶜL_h, ᶜL_v) = p.precomputed
+    c_smag = CAP.c_smag(p.params)
+    grav = CAP.grav(p.params)
+    Pr_t = CAP.Prandtl_number_0(CAP.turbconv_params(p.params))
+    thermo_params = CAP.thermodynamics_params(p.params)
+    (; ᶜtemp_scalar, ᶜtemp_scalar_2) = p.scratch
+    
     # Strain rate tensor
-    ᶜS = @. ᶜtemp_strain = (∇ᶜu_uvw + adjoint(∇ᶜu_uvw)) / 2
-    ᶠS = @. ᶠtemp_strain = (∇ᶠu_uvw + adjoint(∇ᶠu_uvw)) / 2
+    ᶜS = compute_strain_rate_center(ᶠu)
 
     # Stratification correction
-    ᶜθ_v = @. ᶜtemp_scalar = TD.virtual_pottemp(thermo_params, ᶜts)
+    ᶜθ_v = @. lazy(TD.virtual_pottemp(thermo_params, ᶜts))
     ᶜ∇ᵥθ = @. ᶜtemp_scalar_2 = Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜθ_v))).components.data.:1
     ᶜN² = @. ᶜtemp_scalar = grav / ᶜθ_v * ᶜ∇ᵥθ
-    ᶜS_norm = @. ᶜtemp_scalar_2 = √(2 * CA.norm_sqr(ᶜS))
+    ᶜS_norm = @. ᶜtemp_scalar_2 = √(2 * norm_sqr(ᶜS))
 
     ᶜRi = @. ᶜtemp_scalar = ᶜN² / (ᶜS_norm^2 + eps(FT))  # Ri = N² / |S|²
     ᶜfb = @. ᶜtemp_scalar = ifelse(ᶜRi ≤ 0, 1, max(0, 1 - ᶜRi / Pr_t)^(1 / 4))
 
     # filter scale
     h_space = Spaces.horizontal_space(axes(Y.c))
-    Δ_xy = Spaces.node_horizontal_length_scale(h_space)^2 # Δ_x * Δ_y
+    Δ_h = Spaces.node_horizontal_length_scale(h_space)
     ᶜΔ_z = Fields.Δz_field(Y.c)
-    ᶜΔ = @. ᶜtemp_scalar = ∛(Δ_xy * ᶜΔ_z) * ᶜfb
 
-    # Smagorinsky-Lilly eddy viscosity
-    ᶜνₜ = @. ᶜtemp_scalar = c_smag^2 * ᶜΔ^2 * ᶜS_norm
-    ᶠνₜ = @. ᶠtemp_scalar = ᶠinterp(ᶜνₜ)
-
-    # Subgrid-scale momentum flux tensor, `τ = -2 νₜ ∘ S`
-    @. ᶜτ_smag = -2 * ᶜνₜ * ᶜS
-    @. ᶠτ_smag = -2 * ᶠνₜ * ᶠS
-
-    # Turbulent diffusivity
-    @. ᶜD_smag = ᶜνₜ / Pr_t
-    @. ᶠD_smag = ᶠνₜ / Pr_t
+    @. ᶜL_v = c_smag * ᶜΔ_z * ᶜfb
+    @. ᶜL_h = c_smag * Δ_h
 
     nothing
 end
@@ -94,11 +61,28 @@ horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 
 function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly)
-    (; ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶜts) = p.precomputed
+    (; ᶜu, ᶠu, ᶜts, ᶜL_h) = p.precomputed
+    (; ᶜtemp_UVWxUVW, ᶠtemp_UVWxUVW, ᶜtemp_scalar, ᶠtemp_scalar) = p.scratch
     thermo_params = CAP.thermodynamics_params(p.params)
+    Pr_t = CAP.Prandtl_number_0(CAP.turbconv_params(p.params))
 
     ## Momentum tendencies
-    ᶠρ = @. p.scratch.ᶠtemp_scalar = ᶠinterp(Y.c.ρ)
+    ᶜS = compute_strain_rate_center(ᶠu)
+    ᶠS = compute_strain_rate_face(ᶜu)
+    ᶜS_norm = @. lazy(√(2 * norm_sqr(ᶜS)))  # TODO: This should be the UV-components only
+
+    # Smagorinsky eddy viscosity
+    ᶜνₜ_h = @. lazy(ᶜL_h^2 * ᶜS_norm)
+    ᶠνₜ_h = @. lazy(ᶠinterp(ᶜνₜ_h))
+
+    # Turbulent diffusivity
+    ᶜD_smag = @. ᶜtemp_scalar = ᶜνₜ_h / Pr_t
+    
+    # Subgrid-scale momentum flux tensor, `τ = -2 νₜ ∘ S`
+    ᶜτ_smag = @. ᶜtemp_UVWxUVW = -2 * ᶜνₜ_h * ᶜS  # TODO: Lazify once we can mix lazy horizontal & vertical operations
+    ᶠτ_smag = @. ᶠtemp_UVWxUVW = -2 * ᶠνₜ_h * ᶠS
+
+    ᶠρ = ᶠtemp_scalar .= face_density(Y)
     @. Yₜ.c.uₕ -= C12(wdivₕ(Y.c.ρ * ᶜτ_smag) / Y.c.ρ)
     @. Yₜ.f.u₃ -= C3(wdivₕ(ᶠρ * ᶠτ_smag) / ᶠρ)
 
@@ -121,15 +105,12 @@ function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLill
 
 end
 
-import UnrolledUtilities as UU
-
 function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly)
     FT = eltype(Y)
-    (; sfc_temp_C3, ᶠtemp_scalar) = p.scratch
-    (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, sfc_conditions) = p.precomputed
-    (; ρ_flux_uₕ, ρ_flux_h_tot) = sfc_conditions
-    (; ᶜts) = p.precomputed
+    (; ᶜu, ᶠu, ᶜts, ᶜL_v) = p.precomputed
+    (; ᶜtemp_UVWxUVW, ᶠtemp_UVWxUVW, ᶠtemp_scalar, ᶠtemp_scalar_2) = p.scratch
     thermo_params = CAP.thermodynamics_params(p.params)
+    Pr_t = CAP.Prandtl_number_0(CAP.turbconv_params(p.params))
 
     # Define operators
     ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ
@@ -146,9 +127,25 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly)
         bottom = Operators.SetValue(C3(FT(0))),
     )
 
+    ## Momentum tendencies
+    ᶜS = compute_strain_rate_center(ᶠu)
+    ᶠS = compute_strain_rate_face(ᶜu)
+    ᶜS_norm = @. lazy(√(2 * norm_sqr(ᶜS)))  # TODO: This should be the W-component only
+
+    # Smagorinsky eddy viscosity
+    ᶜνₜ_v = @. lazy(ᶜL_v^2 * ᶜS_norm)
+    ᶠνₜ_v = @. lazy(ᶠinterp(ᶜνₜ_v))
+
+    # Subgrid-scale momentum flux tensor, `τ = -2 νₜ ∘ S`
+    ᶜτ_smag = @. ᶜtemp_UVWxUVW = -2 * ᶜνₜ_v * ᶜS
+    ᶠτ_smag = @. ᶠtemp_UVWxUVW = -2 * ᶠνₜ_v * ᶠS
+
+    # Turbulent diffusivity
+    ᶠD_smag = @. ᶠtemp_scalar_2 = ᶠνₜ_v / Pr_t
+
     # Apply to tendencies
     ## Horizontal momentum tendency
-    ᶠρ = @. ᶠtemp_scalar = ᶠinterp(Y.c.ρ)
+    ᶠρ = ᶠtemp_scalar .= face_density(Y)
     @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(ᶠρ * ᶠτ_smag) / Y.c.ρ)
     ## Apply boundary condition for momentum flux
     @. Yₜ.c.uₕ -= ᶜdivᵥ_uₕ(-(FT(0) * ᶠgradᵥ(Y.c.uₕ))) / Y.c.ρ

src/prognostic_equations/remaining_tendency.jl

@@ -348,9 +348,8 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     # NOTE: All ρa tendencies should be applied before calling this function
     pressure_work_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
 
-    sl = p.atmos.smagorinsky_lilly
-    horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, sl)
-    vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, sl)
+    horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, p.atmos.smagorinsky_horizontal)
+    vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, p.atmos.smagorinsky_vertical)
 
     # NOTE: This will zero out all momentum tendencies in the EDMFX advection test, 
     # where velocities do not evolve

src/solver/model_getters.jl

@@ -172,10 +172,10 @@ end
 function get_smagorinsky_model(parsed_args; dir)
     dir = (; h = "horizontal", v = "vertical")[dir]
     smag_dir = parsed_args["smagorinsky_$dir"]
-    smag = parsed_args["smagorinsky"]
+    smag = parsed_args["smagorinsky_lilly"]  # TODO: Remove "lilly" suffix
     @assert smag_dir isa Bool "'smagorinsky_$dir' must be `true` or `false`."
-    @assert smag isa Bool "'smagorinsky' must be `true` or `false`."
-    @assert smag ⊼ smag_dir "Only one of 'smagorinsky' and 'smagorinsky_$dir' can be set to true."
+    @assert smag isa Bool "'smagorinsky_lilly' must be `true` or `false`."
+    @assert smag ⊼ smag_dir "Only one of 'smagorinsky_lilly' and 'smagorinsky_$dir' can be set to true."
     return smag_dir || smag ? SmagorinskyLilly() : nothing
 end
 

src/solver/types.jl

@@ -582,15 +582,16 @@ end
 
 Groups turbulence convection-related models and types.
 """
-Base.@kwdef struct AtmosTurbconv{EDMFX, TCM, SAM, SEDM, SNPM, SVM, SMM, SL}
+Base.@kwdef struct AtmosTurbconv{EDMFX, TCM, SAM, SEDM, SNPM, SVM, SMM, SH, SV}
     edmfx_model::EDMFX = nothing
     turbconv_model::TCM = nothing
     sgs_adv_mode::SAM = nothing
     sgs_entr_detr_mode::SEDM = nothing
     sgs_nh_pressure_mode::SNPM = nothing
     sgs_vertdiff_mode::SVM = nothing
     sgs_mf_mode::SMM = nothing
-    smagorinsky_lilly::SL = nothing
+    smagorinsky_horizontal::SH = nothing
+    smagorinsky_vertical::SV = nothing
 end
 
 """

src/utils/utilities.jl

@@ -87,7 +87,6 @@ compute_kinetic(Y::Fields.FieldVector) = compute_kinetic(Y.c.uₕ, Y.f.u₃)
 Compute the strain_rate at cell centers from velocity at cell faces.
 """
 function compute_strain_rate_center(u::Fields.Field)
-    @assert eltype(u) <: C123
     axis_uvw = Geometry.UVWAxis()
     return @. lazy(
         (
@@ -119,6 +118,18 @@ function compute_strain_rate_face(u::Fields.Field)
     )
 end
 
+"""
+    face_density(Y)
+
+Compute the density at cell faces from the density at cell centers.
+"""
+function face_density(Y)
+    ᶠJ = Fields.local_geometry_field(Y.f).J
+    ᶜJ = Fields.local_geometry_field(Y.c).J
+
+    @. lazy(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ)
+end
+
 """
     g³³_field(space)