unify SGS vertical diffusion with grid mean

Author: sajjadazimi

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-387
+388
 
 # **README**
 #
@@ -32,6 +32,11 @@
 # 3) (optional) leave a link to the buildkite run that prompted this ref counter bump.
 
 #=
+388
+- Unify SGS vertical diffusion with grid mean. Each PrognosticEDMFX
+  subdomain inherits the grid-mean specific tendency for mseⱼ, q_totⱼ,
+  and matching auto-discovered SGS tracers.
+
 387
 - Change some EDMF parameters
 

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -226,6 +226,14 @@ function edmfx_sgs_diffusive_flux_tendency!(
     turbconv_params = CAP.turbconv_params(params)
     (; ᶜu) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
+    n = n_mass_flux_subdomains(turbconv_model)
+    # The SGS-updraft branches below apply the same specific tendency the grid
+    # mean receives to each subdomain scalar (uniform vertical diffusion across
+    # the grid box). Gated on `edmfx_model.vertical_diffusion` so configs can
+    # opt in/out just like the old subdomain-native diffusion did.
+    apply_sgs_updraft =
+        turbconv_model isa PrognosticEDMFX &&
+        p.atmos.edmfx_model.vertical_diffusion isa Val{true}
 
     if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
 
@@ -272,7 +280,13 @@ function edmfx_sgs_diffusive_flux_tendency!(
         ᶜq_vap = @. lazy(
             TD.vapor_specific_humidity(ᶜq_tot_nonneg, ᶜq_liq, ᶜq_ice),
         )
-        @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(
+        # Materialize the total-enthalpy flux divergence once so it can be
+        # applied to the grid-mean ρe_tot and to each subdomain mse using
+        # the same specific-enthalpy tendency dh_gm/dt = dρe_tot/dt / ρ.
+        # Uniform vertical diffusion in the grid box: every subdomain feels
+        # the grid-mean specific tendency.
+        ᶜρe_totₜ_diffusion = p.scratch.ᶜtemp_scalar_2
+        @. ᶜρe_totₜ_diffusion = ᶜdivᵥ_ρe_tot(
             -(
                 ᶠρaK_h * (
                     ᶠgradᵥ(TD.dry_static_energy(thermo_params, ᶜT, ᶜΦ)) +
@@ -285,6 +299,12 @@ function edmfx_sgs_diffusive_flux_tendency!(
                 )
             ),
         )
+        @. Yₜ.c.ρe_tot -= ᶜρe_totₜ_diffusion
+        if apply_sgs_updraft
+            for j in 1:n
+                @. Yₜ.c.sgsʲs.:($$j).mse -= ᶜρe_totₜ_diffusion / Y.c.ρ
+            end
+        end
 
         if use_prognostic_tke(turbconv_model)
             (; ρtke_flux) = p.precomputed
@@ -320,6 +340,12 @@ function edmfx_sgs_diffusive_flux_tendency!(
                 ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(specific(Y.c.ρq_tot, Y.c.ρ))))
             @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion
             @. Yₜ.c.ρ -= ᶜρχₜ_diffusion  # Effect of moisture diffusion on (moist) air mass
+            if apply_sgs_updraft
+                for j in 1:n
+                    @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜρχₜ_diffusion / Y.c.ρ
+                end
+                # The corresponding ρaⱼ dry-mass correction is deliberately neglected.
+            end
         end
 
         α_vert_diff_microphysics = CAP.α_vert_diff_tracer(params)
@@ -356,6 +382,18 @@ function edmfx_sgs_diffusive_flux_tendency!(
                 ),
             )
             @. ᶜρχₜ -= ᶜρχₜ_diffusion
+            # Uniform vertical diffusion: apply the same specific tendency to
+            # the matching subdomain field (e.g. ρq_lcl → q_lcl in updrafts).
+            if apply_sgs_updraft
+                χ_name = specific_tracer_name(ρχ_name)
+                for j in 1:n
+                    if MatrixFields.has_field(Y.c.sgsʲs.:($j), χ_name)
+                        ᶜχⱼₜ =
+                            MatrixFields.get_field(Yₜ.c.sgsʲs.:($j), χ_name)
+                        @. ᶜχⱼₜ -= ᶜρχₜ_diffusion / Y.c.ρ
+                    end
+                end
+            end
         end
 
         # Momentum diffusion

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -54,8 +54,6 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
 
     edmfx_sgs_mass_flux_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
 
-    edmfx_vertical_diffusion_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
-
     edmfx_boundary_condition_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
 
     # NOTE: All ρa tendencies should be applied before calling this function

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -794,9 +794,9 @@ function update_sedimentation_jacobian!(matrix, Y, p, dtγ)
 end
 
 # Center eddy diffusivity and viscosity for the non-EDMF implicit diffusion
-# Jacobians. May write to ᶜtemp_scalar_3. AbstractEDMF configurations return
-# nothing: their grid-mean and updraft diffusion Jacobians both use the
-# face-native ᶠK_h/ᶠK_u/ᶠK_entr from set_face_diffusivities! (see
+# Jacobian. May write to ᶜtemp_scalar_3. AbstractEDMF configurations return
+# nothing: their grid-mean diffusion Jacobian uses the face-native
+# ᶠK_h/ᶠK_u/ᶠK_entr from set_face_diffusivities! (see
 # update_diffusion_jacobian! and update_sgs_diffusion_jacobian!).
 function eddy_diffusivity_coefficients!(Y, p)
     (; vertical_diffusion, smagorinsky_lilly) = p.atmos
@@ -1236,63 +1236,75 @@ end
     update_sgs_diffusion_jacobian!(matrix, Y, p, dtγ, diffusion_flag)
 
 Updates the Jacobian blocks for implicit vertical diffusion of the updraft
-scalars. No-op when diffusion is treated explicitly.
-
-Reuses `ᶜdiffusion_h_matrix` as scratch space, so it must run after
-`update_diffusion_jacobian!`.
+scalars under the unified grid-mean tendency. Under the "uniform diffusion
+in the grid box" tendency in `edmfx_sgs_diffusive_flux_tendency!`, each
+subdomain scalar receives the grid-mean specific tendency which strictly
+depends on grid-mean state, not on χⱼ. The exact linearization would
+put off-diagonal (subdomain-column, grid-mean-column) entries into the
+Jacobian, which the current triangular structure of the ClimaAtmos Jacobian
+does not accommodate. We instead use a `χ ≈ χⱼ` self-linearization (frozen-
+coefficient / diagonal approximation): pretend `ᶜρχₜ_diffusion` depends on χⱼ
+via the same divergence-of-K-gradient operator applied to χⱼ. The residual
+still uses the true (grid-mean-derived) tendency; only the Jacobian is
+approximated. This gives implicit stability for the subdomain scalar
+diffusion without introducing off-triangular Jacobian entries.
 """
 function update_sgs_diffusion_jacobian!(matrix, Y, p, dtγ, diffusion_flag)
     p.atmos.turbconv_model isa PrognosticEDMFX || return nothing
     use_derivative(diffusion_flag) || return nothing
     # Mirror the gate of the tendency this linearizes
-    # (edmfx_vertical_diffusion_tendency!): without it, the updraft scalar
-    # diagonals would carry diffusion terms that have no tendency counterpart.
-    p.atmos.edmfx_model.vertical_diffusion isa Val{true} || return nothing
+    # (`edmfx_sgs_diffusive_flux_tendency!` inside the sgs_diffusive_flux
+    # branch): without it, the updraft scalar diagonals would carry
+    # diffusion terms that have no tendency counterpart.
+    p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} || return nothing
     (; params) = p
-    (; ᶜρʲs) = p.precomputed
     (; ᶜdiffusion_h_matrix) = p.scratch
+    ᶜρ = Y.c.ρ
 
-    α_vert_diff_microphysics = CAP.α_vert_diff_tracer(params)
-    # Face-native ᶠK_h, consistent with edmfx_vertical_diffusion_tendency!
-    # (ᶠK_entr is deliberately excluded there; see that function).
-    (; ᶠK_h) = p.precomputed
-    @. ᶜdiffusion_h_matrix =
-        ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1)) * ᶠK_h) ⋅
-        ᶠgradᵥ_matrix()
-
+    # mseⱼ and q_totⱼ diagonals: same operator, no tracer factor. Uses the
+    # full ρ(K_h + K_entr) diffusion matrix.
     ∂ᶜmseʲ_err_∂ᶜmseʲ =
         matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).mse)]
     ∂ᶜq_totʲ_err_∂ᶜq_totʲ =
         matrix[@name(c.sgsʲs.:(1).q_tot), @name(c.sgsʲs.:(1).q_tot)]
     @. ∂ᶜmseʲ_err_∂ᶜmseʲ +=
-        dtγ * DiagonalMatrixRow(1 / ᶜρʲs.:(1)) ⋅ ᶜdiffusion_h_matrix
+        dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_h_matrix
     @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ +=
-        dtγ * DiagonalMatrixRow(1 / ᶜρʲs.:(1)) ⋅ ᶜdiffusion_h_matrix
+        dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_h_matrix
 
+    # Auto-discovered SGS tracers. Sedimenting microphysics species are
+    # diffused with `α·K_h + K_entr` (matching the tendency); passive
+    # tracers keep the full `K_h + K_entr` — so the passive branch reuses
+    # `ᶜdiffusion_h_matrix` directly and the sedimenting branch rebuilds the
+    # α-scaled matrix in scratch.
     if p.atmos.microphysics_model isa Union{
         NonEquilibriumMicrophysics1M,
         NonEquilibriumMicrophysics2M,
     }
+        α_vert_diff_microphysics = CAP.α_vert_diff_tracer(params)
+        (; ᶠK_h, ᶠK_entr) = p.precomputed
+        ᶜsgs_tracer_diffusion_matrix = p.scratch.ᶜtridiagonal_matrix_scalar
+        @. ᶜsgs_tracer_diffusion_matrix =
+            ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(
+                ᶠinterp(ᶜρ) * (α_vert_diff_microphysics * ᶠK_h + ᶠK_entr),
+            ) ⋅ ᶠgradᵥ_matrix()
         MatrixFields.unrolled_foreach(
             sedimenting_sgs_tracer_names(Y),
         ) do χ_name
             χ_state_name = sgs_state_name(χ_name)
             ∂ᶜχʲ_err_∂ᶜχʲ = matrix[χ_state_name, χ_state_name]
             @. ∂ᶜχʲ_err_∂ᶜχʲ +=
-                dtγ * α_vert_diff_microphysics *
-                DiagonalMatrixRow(1 / ᶜρʲs.:(1)) ⋅
-                ᶜdiffusion_h_matrix
+                dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅
+                ᶜsgs_tracer_diffusion_matrix
         end
     end
 
-    # Passive SGS tracers are diffused with the unscaled K_h (see
-    # edmfx_vertical_diffusion_tendency!); their diagonals are initialized by
-    # update_sgs_advection_jacobian!, so the diffusion term is accumulated.
+    # Passive SGS tracers: unscaled ρ(K_h + K_entr), same as mseⱼ / q_totⱼ.
     MatrixFields.unrolled_foreach(passive_sgs_tracer_names(Y)) do χ_name
         χ_state_name = sgs_state_name(χ_name)
         ∂ᶜχʲ_err_∂ᶜχʲ = matrix[χ_state_name, χ_state_name]
         @. ∂ᶜχʲ_err_∂ᶜχʲ +=
-            dtγ * DiagonalMatrixRow(1 / ᶜρʲs.:(1)) ⋅ ᶜdiffusion_h_matrix
+            dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_h_matrix
     end
     return nothing
 end

src/prognostic_equations/mass_flux_closures.jl

@@ -162,77 +162,6 @@ layers (`⟨w'b'⟩_s ≤ 0`).
     return a_s * ρ * w_star
 end
 
-edmfx_vertical_diffusion_tendency!(Yₜ, Y, p, t, turbconv_model) = nothing
-
-function edmfx_vertical_diffusion_tendency!(
-    Yₜ,
-    Y,
-    p,
-    t,
-    turbconv_model::PrognosticEDMFX,
-)
-    if p.atmos.edmfx_model.vertical_diffusion isa Val{true}
-        (; params) = p
-        (; ᶜρʲs) = p.precomputed
-        FT = eltype(p.params)
-        n = n_mass_flux_subdomains(turbconv_model)
-        ᶜdivᵥ_mse = Operators.DivergenceF2C(
-            top = Operators.SetValue(C3(0)),
-            bottom = Operators.SetValue(C3(0)),
-        )
-        ᶜdivᵥ_q_tot = Operators.DivergenceF2C(
-            top = Operators.SetValue(C3(0)),
-            bottom = Operators.SetValue(C3(0)),
-        )
-
-        # Updraft internal diffusion uses the same face-native environment
-        # diffusivity ᶠK_h as the grid-mean diffusion (see
-        # set_face_diffusivities! and edmfx_sgs_diffusive_flux_tendency!):
-        # the SGS turbulence that stirs updraft interiors is the same
-        # environment turbulence, and the face-native, interface-aware
-        # evaluation collapses the flux at faces bordering quiescent,
-        # strongly stratified air without interpolation. ᶠK_entr is
-        # deliberately not added here: it represents grid-mean interfacial
-        # entrainment, which for the updrafts is carried by the
-        # entrainment/detrainment closures.
-        (; ᶠK_h) = p.precomputed
-        for j in 1:n
-            ᶜρʲ = ᶜρʲs.:($j)
-            ᶜmseʲ = Y.c.sgsʲs.:($j).mse
-            ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot
-            # Note: For this and other diffusive tendencies, we should use ρaʲ instead of ρʲ,
-            # but it causes stability issues when ρaʲ is small
-            @. Yₜ.c.sgsʲs.:($$j).mse -=
-                ᶜdivᵥ_mse(-(ᶠinterp(ᶜρʲ) * ᶠK_h * ᶠgradᵥ(ᶜmseʲ))) / ᶜρʲ
-            @. Yₜ.c.sgsʲs.:($$j).q_tot -=
-                ᶜdivᵥ_q_tot(-(ᶠinterp(ᶜρʲ) * ᶠK_h * ᶠgradᵥ(ᶜq_totʲ))) / ᶜρʲ
-        end
-
-        if !isempty(sgs_tracer_names(Y))
-            α_vert_diff_microphysics = CAP.α_vert_diff_tracer(params)
-            ᶜρʲ = ᶜρʲs.:(1)
-            ᶜdivᵥ_q = Operators.DivergenceF2C(
-                top = Operators.SetValue(C3(FT(0))),
-                bottom = Operators.SetValue(C3(FT(0))),
-            )
-            # Sedimenting microphysics species are diffused with
-            # α_vert_diff_tracer * K_h, passive tracers with the unscaled K_h,
-            # matching the grid-mean tracer diffusion and the implicit
-            # Jacobian (update_sgs_diffusion_jacobian!).
-            for χ_name in sgs_tracer_names(Y)
-                α =
-                    χ_name in sgs_sedimenting_tracer_candidates ?
-                    α_vert_diff_microphysics :
-                    one(α_vert_diff_microphysics)
-                ᶜχʲ = MatrixFields.get_field(Y.c.sgsʲs.:(1), χ_name)
-                ᶜχʲₜ = MatrixFields.get_field(Yₜ.c.sgsʲs.:(1), χ_name)
-                @. ᶜχʲₜ -=
-                    ᶜdivᵥ_q(-(ᶠinterp(ᶜρʲ) * ᶠK_h * α * ᶠgradᵥ(ᶜχʲ))) / ᶜρʲ
-            end
-        end
-    end
-end
-
 # Private helper: clips grid-mean condensate tracers to non-negative values and
 # rescales the condensate sum so it cannot exceed the available total moisture.
 function enforce_grid_mean_microphysics_constraints!(Y, p, t)