rft: separate vertical and full strain

Author: haakon-e

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -1376,8 +1376,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     # TODO: Currently the shear production only includes vertical gradients
     ᶠu = p.scratch.ᶠtemp_C123
     @. ᶠu = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³)
-    ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
-    ᶜstrain_rate .= compute_strain_rate_center(ᶠu)
+    ᶜstrain_rate = compute_strain_rate_center_vertical(ᶠu)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ρatke_flux_values = Fields.field_values(ρatke_flux)

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -451,8 +451,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     # TODO: Currently the shear production only includes vertical gradients
     ᶠu = p.scratch.ᶠtemp_C123
     @. ᶠu = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³)
-    ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
-    ᶜstrain_rate .= compute_strain_rate_center(ᶠu)
+    ᶜstrain_rate = compute_strain_rate_center_vertical(ᶠu)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ρatke_flux_values = Fields.field_values(ρatke_flux)

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -499,8 +499,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
         end
 
         # Momentum diffusion
-        ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
-        ᶠstrain_rate .= compute_strain_rate_face(ᶜu⁰)
+        ᶠstrain_rate = compute_strain_rate_face_vertical(ᶜu⁰)
         @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ)
     end
     return nothing
@@ -625,8 +624,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
         end
 
         # Momentum diffusion
-        ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
-        ᶠstrain_rate .= compute_strain_rate_face(ᶜu)
+        ᶠstrain_rate = compute_strain_rate_face_vertical(ᶜu)
         @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ)
     end
 

src/prognostic_equations/gm_sgs_closures.jl

@@ -84,8 +84,7 @@ NVTX.@annotate function compute_gm_mixing_length(Y, p)
     # TODO: move strain rate calculation to separate function
     ᶠu = p.scratch.ᶠtemp_C123
     @. ᶠu = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³)
-    ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
-    ᶜstrain_rate .= compute_strain_rate_center(ᶠu)
+    ᶜstrain_rate = compute_strain_rate_center_vertical(ᶠu)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_2

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -104,8 +104,7 @@ function vertical_diffusion_boundary_layer_tendency!(
     end
 
     if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion)
-        ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
-        ᶠstrain_rate .= compute_strain_rate_face(ᶜu)
+        ᶠstrain_rate = compute_strain_rate_face_vertical(ᶜu)
         @. Yₜ.c.uₕ -= C12(
             ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠstrain_rate) / Y.c.ρ,
         ) # assumes ᶜK_u = ᶜK_h

src/utils/utilities.jl

@@ -82,43 +82,116 @@ state.
 compute_kinetic(Y::Fields.FieldVector) = compute_kinetic(Y.c.uₕ, Y.f.u₃)
 
 """
-    ϵ .= compute_strain_rate_center(u::Field)
+    ϵ .= compute_strain_rate_center_vertical(ᶠu)
 
-Compute the strain_rate at cell centers from velocity at cell faces.
+Compute the strain rate at cell centers from velocity at cell faces, with vertical gradients only.
+
+Returns a lazy representation of the strain rate tensor.
 """
-function compute_strain_rate_center(u::Fields.Field)
-    @assert eltype(u) <: C123
+function compute_strain_rate_center_vertical(ᶠu)
     axis_uvw = Geometry.UVWAxis()
     return @. lazy(
         (
-            Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(u))) +
-            adjoint(Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(u))))
+            Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(ᶠu))) +
+            adjoint(Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(ᶠu))))
         ) / 2,
     )
 end
 
 """
-    ϵ .= compute_strain_rate_face(u::Field)
+    ϵ .= compute_strain_rate_face_vertical(ᶜu::Field)
+
+Compute the strain rate at cell faces from velocity at cell centers, with vertical gradients only.
 
-Compute the strain_rate at cell faces from velocity at cell centers.
+Returns a lazy representation of the strain rate tensor.
 """
-function compute_strain_rate_face(u::Fields.Field)
-    @assert eltype(u) <: C123
-    ∇ᵥuvw_boundary =
-        Geometry.outer(Geometry.WVector(0), Geometry.UVWVector(0, 0, 0))
-    ᶠgradᵥ = Operators.GradientC2F(
-        bottom = Operators.SetGradient(∇ᵥuvw_boundary),
-        top = Operators.SetGradient(∇ᵥuvw_boundary),
-    )
+function compute_strain_rate_face_vertical(ᶜu)
+    ∇ᵥuvw_boundary = Geometry.outer(Geometry.WVector(0), Geometry.UVWVector(0, 0, 0))
+    ∇bc = Operators.SetGradient(∇ᵥuvw_boundary)
+    ᶠgradᵥ = Operators.GradientC2F(bottom = ∇bc, top = ∇bc)
     axis_uvw = Geometry.UVWAxis()
     return @. lazy(
         (
-            Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u))) +
-            adjoint(Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u))))
+            Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(ᶜu))) +
+            adjoint(Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(ᶜu))))
         ) / 2,
     )
 end
 
+"""
+    compute_strain_rate_center_full!(ᶜε, ᶜu, ᶠu)
+
+Compute the full strain rate tensor at cell centers from velocity
+
+# Arguments
+ - `ᶜε`: Preallocated `UVW x UVW` tensor field for the strain rate at cell centers
+ - `ᶜu, ᶠu`: Velocity field at cell centers and faces, respectively.
+    Both reconstructions are needed to compute the full strain rate tensor.
+
+See also [`compute_strain_rate_face_full!`](@ref) for the analogous calculation on cell faces.
+
+# Notes:
+- it is recommended to use `ᶠu` and `ᶜu` as computed by 
+    [`set_velocity_quantities!`](@ref) and [`set_implicit_precomputed_quantities_part1!`](@ref)
+- Because the computation involves both vertical and horizontal gradients, this
+    calculation cannot be lazified (for now). It requires a pre-allocated output field.
+"""
+function compute_strain_rate_center_full!(ᶜε, ᶜu, ᶠu)
+    axis_uvw = (Geometry.UVWAxis(),)
+    @. ᶜε = Geometry.project(axis_uvw, ᶜgradᵥ(UVW(ᶠu)))  # vertical component
+    @. ᶜε += Geometry.project(axis_uvw, gradₕ(UVW(ᶜu)))  # horizontal component
+    @. ᶜε = (ᶜε + adjoint(ᶜε)) / 2
+    return ᶜε
+end
+
+"""
+    compute_strain_rate_face_full!(ᶠε, ᶜu, ᶠu)
+
+Compute the full strain rate tensor at cell faces from velocity
+
+# Arguments
+ - `ᶠε`: Preallocated `UVW x UVW` tensor field for the strain rate at cell centers
+ - `ᶜu, ᶠu`: Velocity field at cell centers and faces, respectively.
+    Both reconstructions are needed to compute the full strain rate tensor.
+
+See also [`compute_strain_rate_center_full!`](@ref) for the analogous calculation on cell centers.
+
+# Notes:
+- it is recommended to use `ᶠu` and `ᶜu` as computed by 
+    [`set_velocity_quantities!`](@ref) and [`set_implicit_precomputed_quantities_part1!`](@ref)
+- Because the computation involves both vertical and horizontal gradients, this
+    calculation cannot be lazified (for now). It requires a pre-allocated output field.
+- The calculation assumes no-flux boundary conditions
+"""
+function compute_strain_rate_face_full!(ᶠε, ᶜu, ᶠu)
+    ∇ᵥuvw_boundary = Geometry.outer(Geometry.WVector(0), UVW(0, 0, 0))
+    ∇bc = Operators.SetGradient(∇ᵥuvw_boundary)
+    ᶠgradᵥ = Operators.GradientC2F(bottom = ∇bc, top = ∇bc)
+    axis_uvw = (Geometry.UVWAxis(),)
+    @. ᶠε = Geometry.project(axis_uvw, ᶠgradᵥ(UVW(ᶜu)))  # vertical component
+    @. ᶠε += Geometry.project(axis_uvw, gradₕ(UVW(ᶠu)))  # horizontal component
+    @. ᶠε = (ᶠε + adjoint(ᶠε)) / 2
+    return ᶠε
+end
+
+"""
+    strain_rate_norm(S, axis = Geometry.UVWAxis())
+
+Return a lazy representation of the strain rate norm `|S| = √(2 ∘ S : S)`
+
+If `axis` is provided, project the strain rate tensor `S` onto the specified axis 
+before computing the norm. 
+
+For example, 
+- `axis = Geometry.UVAxis()` computes the horizontal strain rate norm, while 
+- `axis = Geometry.WAxis()` computes the vertical strain rate norm.
+"""
+function strain_rate_norm(S, axis = Geometry.UVWAxis())
+    S_proj = @. lazy(Geometry.project((axis,), S, (axis,)))
+    S_norm = @. lazy(√(2 * norm_sqr(S_proj)))
+    return S_norm
+end
+
 """
     g³³_field(space)
 

test/utilities.jl

@@ -89,7 +89,7 @@ end
 end
 
 @testset "compute_strain_rate (c.f. analytical function)" begin
-    # Test compute_strain_rate_face
+    # Test compute_strain_rate_face_vertical
     (; helem, cent_space, face_space) = get_cartesian_spaces()
     ccoords, fcoords = get_coords(cent_space, face_space)
     FT = eltype(ccoords.x)
@@ -118,8 +118,8 @@ end
     ᶠu = @. UVW(Geometry.UVector(ᶠu)) +
        UVW(Geometry.VVector(ᶠv)) +
        UVW(Geometry.WVector(ᶠw))
-    ᶜϵ .= CA.compute_strain_rate_center(Geometry.Covariant123Vector.(ᶠu))
-    ᶠϵ .= CA.compute_strain_rate_face(Geometry.Covariant123Vector.(ᶜu))
+    ᶜϵ .= CA.compute_strain_rate_center_vertical(Geometry.Covariant123Vector.(ᶠu))
+    ᶠϵ .= CA.compute_strain_rate_face_vertical(Geometry.Covariant123Vector.(ᶜu))
 
     # Center valued strain rate
     @test ᶜϵ.components.data.:1 == ᶜϵ.components.data.:1 .* FT(0)