Merge pull request #4073 from CliMA/he/rft-vert-full-strain

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 zero vertical gradient 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(sqrt(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)
@@ -163,6 +163,96 @@ end
     end
 end
 
+@testset "compute_full_strain_rate (consistency and symmetry)" begin
+    (; helem, cent_space, face_space) = get_cartesian_spaces()
+    ccoords, fcoords = get_coords(cent_space, face_space)
+    FT = eltype(ccoords.x)
+    UVW = Geometry.UVWVector
+    C123 = Geometry.Covariant123Vector
+    UVec = Geometry.UVector
+    VVec = Geometry.VVector
+    WVec = Geometry.WVector
+    ᶜgradᵥ = Operators.GradientF2C()
+    gradₕ = Operators.Gradient()
+
+
+    # Alloc scratch space for 3x3 tensor fields
+    u₀ = UVW(FT(0), FT(0), FT(0))
+    ᶜε = Fields.Field(typeof(u₀ * u₀'), cent_space)
+    ᶠε = Fields.Field(typeof(u₀ * u₀'), face_space)
+    ᶜε_ref = Fields.Field(typeof(u₀ * u₀'), cent_space)
+    ᶠε_ref = Fields.Field(typeof(u₀ * u₀'), face_space)
+
+    # Velocity fields
+    u, v, w = taylor_green_ic(ccoords)
+    ᶠu, ᶠv, ᶠw = taylor_green_ic(fcoords)
+    ᶜu = @. UVW(UVec(u)) + UVW(VVec(v)) + UVW(WVec(w))
+    ᶠu_vec = @. UVW(UVec(ᶠu)) + UVW(VVec(ᶠv)) + UVW(WVec(ᶠw))
+
+    # Compute using API under test
+    CA.compute_strain_rate_center_full!(ᶜε, ᶜu, ᶠu_vec)
+    CA.compute_strain_rate_face_full!(ᶠε, ᶜu, ᶠu_vec)
+
+    # Build reference tensors explicitly (projection + symmetrization)
+    axis_uvw = (Geometry.UVWAxis(),)
+    # Center reference
+    @. ᶜε_ref = Geometry.project(axis_uvw, ᶜgradᵥ(ᶠu_vec))
+    @. ᶜε_ref += Geometry.project(axis_uvw, gradₕ(ᶜu))
+    @. ᶜε_ref = (ᶜε_ref + adjoint(ᶜε_ref)) / 2
+
+    # Face reference (construct vertical gradient with the same BCs)
+    ∇ᵥuvw_boundary = Geometry.outer(WVec(0), UVW(0, 0, 0))
+    ∇bc = Operators.SetGradient(∇ᵥuvw_boundary)
+    ᶠgradᵥ = Operators.GradientC2F(bottom = ∇bc, top = ∇bc)
+    @. ᶠε_ref = Geometry.project(axis_uvw, ᶠgradᵥ(ᶜu))
+    @. ᶠε_ref += Geometry.project(axis_uvw, gradₕ(ᶠu_vec))
+    @. ᶠε_ref = (ᶠε_ref + adjoint(ᶠε_ref)) / 2
+
+    # Symmetry checks (independent of reference)
+    @test ᶜε.components.data.:2 == ᶜε.components.data.:4
+    @test ᶜε.components.data.:3 == ᶜε.components.data.:7
+    @test ᶜε.components.data.:6 == ᶜε.components.data.:8
+    @test ᶠε.components.data.:2 == ᶠε.components.data.:4
+    @test ᶠε.components.data.:3 == ᶠε.components.data.:7
+    @test ᶠε.components.data.:6 == ᶠε.components.data.:8
+
+    # Consistency with reference (component-wise, tight tolerance)
+    tol = sqrt(eps(FT))
+    @test maximum(abs, ᶜε.components.data.:1 .- ᶜε_ref.components.data.:1) < tol
+    @test maximum(abs, ᶜε.components.data.:2 .- ᶜε_ref.components.data.:2) < tol
+    @test maximum(abs, ᶜε.components.data.:3 .- ᶜε_ref.components.data.:3) < tol
+    @test maximum(abs, ᶜε.components.data.:4 .- ᶜε_ref.components.data.:4) < tol
+    @test maximum(abs, ᶜε.components.data.:5 .- ᶜε_ref.components.data.:5) < tol
+    @test maximum(abs, ᶜε.components.data.:6 .- ᶜε_ref.components.data.:6) < tol
+    @test maximum(abs, ᶜε.components.data.:7 .- ᶜε_ref.components.data.:7) < tol
+    @test maximum(abs, ᶜε.components.data.:8 .- ᶜε_ref.components.data.:8) < tol
+    @test maximum(abs, ᶜε.components.data.:9 .- ᶜε_ref.components.data.:9) < tol
+
+    @test maximum(abs, ᶠε.components.data.:1 .- ᶠε_ref.components.data.:1) < tol
+    @test maximum(abs, ᶠε.components.data.:2 .- ᶠε_ref.components.data.:2) < tol
+    @test maximum(abs, ᶠε.components.data.:3 .- ᶠε_ref.components.data.:3) < tol
+    @test maximum(abs, ᶠε.components.data.:4 .- ᶠε_ref.components.data.:4) < tol
+    @test maximum(abs, ᶠε.components.data.:5 .- ᶠε_ref.components.data.:5) < tol
+    @test maximum(abs, ᶠε.components.data.:6 .- ᶠε_ref.components.data.:6) < tol
+    @test maximum(abs, ᶠε.components.data.:7 .- ᶠε_ref.components.data.:7) < tol
+    @test maximum(abs, ᶠε.components.data.:8 .- ᶠε_ref.components.data.:8) < tol
+    @test maximum(abs, ᶠε.components.data.:9 .- ᶠε_ref.components.data.:9) < tol
+
+    # Boundary behavior (face): match reference at first and last vertical levels per element
+    for elem_id in 1:helem
+        @test maximum(
+            abs,
+            Fields.field_values(Fields.slab(ᶠε.components.data.:3, 1, elem_id)) .-
+            Fields.field_values(Fields.slab(ᶠε_ref.components.data.:3, 1, elem_id)),
+        ) < tol
+        @test maximum(
+            abs,
+            Fields.field_values(Fields.slab(ᶠε.components.data.:3, 11, elem_id)) .-
+            Fields.field_values(Fields.slab(ᶠε_ref.components.data.:3, 11, elem_id)),
+        ) < tol
+    end
+end
+
 @testset "horizontal integral at boundary" begin
     # Test both `horizontal_integral_at_boundary` methods
     (; cent_space, face_space) = get_cartesian_spaces()