Merge pull request #4014 from CliMA/zs/implicit_updraft_diffusion

szy21 · web-flow · commit cf4677e3ca2f · 2025-09-15T21:56:05.000Z
make updraft diffusion implicit
diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
@@ -353,6 +353,9 @@ implicit_sgs_nh_pressure:
   `implicit_sgs_advection` is set to true. This flag only controls whether the drag term in the pressure closure is treated implicitly. 
   The buoyancy term is always treated implicitly. [`false` (default), `true`]"
   value: false
+implicit_sgs_vertdiff:
+  help: "Whether to treat the subgrid-scale vertical diffusion tendency implicitly [`false` (default), `true`]. Setting it to true only works if implicit_sgs_advection is set to true."
+  value: false
 implicit_sgs_mass_flux:
   help: "Whether to treat the subgrid-scale mass flux tendency implicitly or explicitly in grid-mean equations. Currently updraft only with Jacobian terms 0. [`false` (default), `true`]. Setting it to true only works if both implicit_sgs_advection and implicit_diffusion are set to true."
   value: false
diff --git a/config/longrun_configs/longrun_aquaplanet_allsky_progedmf_0M.yml b/config/longrun_configs/longrun_aquaplanet_allsky_progedmf_0M.yml
@@ -2,6 +2,7 @@ implicit_sgs_advection: true
 implicit_sgs_mass_flux: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 max_newton_iters_ode: 3
 dt: "30secs"
 t_end: "60days"
@@ -17,6 +18,7 @@ edmfx_upwinding: "first_order"
 edmfx_entr_model: "Generalized"
 edmfx_detr_model: "Generalized"
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
diff --git a/config/model_configs/prognostic_edmfx_aquaplanet.yml b/config/model_configs/prognostic_edmfx_aquaplanet.yml
@@ -20,6 +20,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/model_configs/prognostic_edmfx_aquaplanet_dense_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_dense_autodiff.yml
@@ -22,6 +22,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 use_dense_jacobian: true
 update_jacobian_every: dt
diff --git a/config/model_configs/prognostic_edmfx_aquaplanet_gpu.yml b/config/model_configs/prognostic_edmfx_aquaplanet_gpu.yml
@@ -22,6 +22,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/model_configs/prognostic_edmfx_aquaplanet_gpu_dense_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_gpu_dense_autodiff.yml
@@ -22,6 +22,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 use_dense_jacobian: true
 update_jacobian_every: dt
diff --git a/config/model_configs/prognostic_edmfx_aquaplanet_gpu_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_gpu_sparse_autodiff.yml
@@ -22,6 +22,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 use_auto_jacobian: true
 update_jacobian_every: dt
diff --git a/config/model_configs/prognostic_edmfx_aquaplanet_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_sparse_autodiff.yml
@@ -22,6 +22,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 use_auto_jacobian: true
 update_jacobian_every: dt
diff --git a/config/model_configs/prognostic_edmfx_bomex_column_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_bomex_column_sparse_autodiff.yml
@@ -8,6 +8,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml b/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml
@@ -8,6 +8,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/model_configs/prognostic_edmfx_diurnal_scm_imp.yml b/config/model_configs/prognostic_edmfx_diurnal_scm_imp.yml
@@ -14,6 +14,8 @@ implicit_sgs_advection: true
 implicit_sgs_mass_flux: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+# unstable with implicit vertical diffusion
+implicit_sgs_vertdiff: false
 
 approximate_linear_solve_iters: 2
 edmfx_upwinding: first_order
@@ -26,7 +28,7 @@ edmfx_detr_model: "PiGroups"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
-edmfx_vertical_diffusion: true
+edmfx_vertical_diffusion: false
 edmfx_filter: true
 prognostic_tke: true
 precip_model: "0M"
diff --git a/config/model_configs/prognostic_edmfx_dycoms_rf02_column_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_dycoms_rf02_column_sparse_autodiff.yml
@@ -8,6 +8,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/model_configs/prognostic_edmfx_gabls_column_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_gabls_column_sparse_autodiff.yml
@@ -6,6 +6,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/model_configs/prognostic_edmfx_rico_implicit_column.yml b/config/model_configs/prognostic_edmfx_rico_implicit_column.yml
@@ -7,6 +7,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_nh_pressure: true
+implicit_sgs_vertdiff: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 edmfx_upwinding: first_order
diff --git a/config/model_configs/prognostic_edmfx_trmm_column_0M.yml b/config/model_configs/prognostic_edmfx_trmm_column_0M.yml
@@ -6,6 +6,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/model_configs/prognostic_edmfx_trmm_column_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_trmm_column_sparse_autodiff.yml
@@ -6,6 +6,7 @@ implicit_diffusion: true
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
diff --git a/config/perf_configs/bm_aquaplanet_progedmf.yml b/config/perf_configs/bm_aquaplanet_progedmf.yml
@@ -7,6 +7,7 @@ log_progress: false
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 moist: equil
 surface_setup: DefaultMoninObukhov
diff --git a/config/perf_configs/bm_aquaplanet_progedmf_dense_autodiff.yml b/config/perf_configs/bm_aquaplanet_progedmf_dense_autodiff.yml
@@ -8,6 +8,7 @@ approximate_linear_solve_iters: 1
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 use_dense_jacobian: true
 update_jacobian_every: dt
diff --git a/config/perf_configs/bm_aquaplanet_progedmf_sparse_autodiff.yml b/config/perf_configs/bm_aquaplanet_progedmf_sparse_autodiff.yml
@@ -8,6 +8,7 @@ approximate_linear_solve_iters: 1
 implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_sgs_nh_pressure: true
+implicit_sgs_vertdiff: true
 implicit_sgs_mass_flux: true
 use_auto_jacobian: true
 update_jacobian_every: dt
diff --git a/reproducibility_tests/ref_counter.jl b/reproducibility_tests/ref_counter.jl
@@ -1,4 +1,4 @@
-262
+263
 
 # **README**
 #
@@ -20,6 +20,9 @@
 
 
 #=
+263
+- Make vertical diffusion of updrafts implicit
+
 262
 - Add vertical diffusion to prognostic EDMF updrafts
 
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -55,6 +55,10 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
         edmfx_nh_pressure_drag_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     end
 
+    if p.atmos.sgs_vertdiff_mode == Implicit()
+        edmfx_vertical_diffusion_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    end
+
     # NOTE: All ρa tendencies should be applied before calling this function
     pressure_work_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
 
diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl
@@ -22,6 +22,7 @@ use_derivative(::IgnoreDerivative) = false
         sgs_entr_detr_flag,
         sgs_mass_flux_flag,
         sgs_nh_pressure_flag,
+        sgs_vertdiff_flag,
         approximate_solve_iters,
     )
 
@@ -45,16 +46,19 @@ solver. Certain groups of derivatives can be toggled on or off by setting their
   subgrid-scale mass flux tendency should be computed
 - `sgs_nh_pressure_flag::DerivativeFlag`: whether the derivatives of the
   subgrid-scale non-hydrostatic pressure drag tendency should be computed
+- `sgs_vertdiff_flag::DerivativeFlag`: whether the derivatives of the
+  subgrid-scale vertical diffusion tendency should be computed
 - `approximate_solve_iters::Int`: number of iterations to take for the
   approximate linear solve required when the `diffusion_flag` is `UseDerivative`
 """
-struct ManualSparseJacobian{F1, F2, F3, F4, F5, F6} <: SparseJacobian
+struct ManualSparseJacobian{F1, F2, F3, F4, F5, F6, F7} <: SparseJacobian
     topography_flag::F1
     diffusion_flag::F2
     sgs_advection_flag::F3
     sgs_entr_detr_flag::F4
     sgs_mass_flux_flag::F5
     sgs_nh_pressure_flag::F6
+    sgs_vertdiff_flag::F7
     approximate_solve_iters::Int
 end
 
@@ -336,6 +340,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
         sgs_entr_detr_flag,
         sgs_nh_pressure_flag,
         sgs_mass_flux_flag,
+        sgs_vertdiff_flag,
     ) = alg
     (; matrix) = cache
     (; params) = p
@@ -976,6 +981,49 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
                     DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - (I_u₃,)
             end
 
+            # vertical diffusion of updrafts
+            if use_derivative(sgs_vertdiff_flag)
+                α_vert_diff_tracer = CAP.α_vert_diff_tracer(params)
+                @. ᶜdiffusion_h_matrix =
+                    ᶜadvdivᵥ_matrix() ⋅
+                    DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1)) * ᶠinterp(ᶜK_h)) ⋅
+                    ᶠgradᵥ_matrix()
+                @. ᶜdiffusion_h_matrix_scaled =
+                    ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(
+                        ᶠinterp(ᶜρʲs.:(1)) * ᶠinterp(α_vert_diff_tracer * ᶜK_h),
+                    ) ⋅ ᶠgradᵥ_matrix()
+
+                @. ∂ᶜmseʲ_err_∂ᶜmseʲ +=
+                    dtγ * DiagonalMatrixRow(1 / ᶜρʲs.:(1)) ⋅ ᶜdiffusion_h_matrix
+                @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ +=
+                    dtγ * DiagonalMatrixRow(1 / ᶜρʲs.:(1)) ⋅ ᶜdiffusion_h_matrix
+                if p.atmos.moisture_model isa NonEquilMoistModel && (
+                    p.atmos.microphysics_model isa Microphysics1Moment ||
+                    p.atmos.microphysics_model isa Microphysics2Moment
+                )
+                    sgs_microphysics_tracers = (
+                        (@name(c.sgsʲs.:(1).q_liq)),
+                        (@name(c.sgsʲs.:(1).q_ice)),
+                        (@name(c.sgsʲs.:(1).q_rai)),
+                        (@name(c.sgsʲs.:(1).q_sno)),
+                    )
+                    MatrixFields.unrolled_foreach(
+                        sgs_microphysics_tracers,
+                    ) do (qʲ_name)
+                        MatrixFields.has_field(Y, qʲ_name) || return
+                        ᶜtridiagonal_matrix_scalar = ifelse(
+                            qʲ_name in
+                            (@name(c.sgsʲs.:(1).q_rai), @name(c.sgsʲs.:(1).q_snow)),
+                            ᶜdiffusion_h_matrix_scaled,
+                            ᶜdiffusion_h_matrix,
+                        )
+                        ∂ᶜqʲ_err_∂ᶜqʲ = matrix[qʲ_name, qʲ_name]
+                        @. ∂ᶜqʲ_err_∂ᶜqʲ +=
+                            dtγ * DiagonalMatrixRow(1 / ᶜρʲs.:(1)) ⋅
+                            ᶜtridiagonal_matrix_scalar
+                    end
+                end
+            end
             # entrainment and detrainment (rates are treated explicitly)
             if use_derivative(sgs_entr_detr_flag)
                 (; ᶜentrʲs, ᶜdetrʲs, ᶜturb_entrʲs) = p.precomputed
diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl
@@ -290,7 +290,9 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     if p.atmos.sgs_nh_pressure_mode == Explicit()
         edmfx_nh_pressure_drag_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     end
-    edmfx_vertical_diffusion_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    if p.atmos.sgs_vertdiff_mode == Explicit()
+        edmfx_vertical_diffusion_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    end
     edmfx_filter_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     edmfx_tke_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
 
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
@@ -86,6 +86,9 @@ function get_atmos(config::AtmosConfig, params)
     implicit_sgs_nh_pressure = parsed_args["implicit_sgs_nh_pressure"]
     @assert implicit_sgs_nh_pressure in (true, false)
 
+    implicit_sgs_vertdiff = parsed_args["implicit_sgs_vertdiff"]
+    @assert implicit_sgs_vertdiff in (true, false)
+
     implicit_sgs_mass_flux = parsed_args["implicit_sgs_mass_flux"]
     @assert implicit_sgs_mass_flux in (true, false)
 
@@ -136,6 +139,7 @@ function get_atmos(config::AtmosConfig, params)
         sgs_entr_detr_mode = implicit_sgs_entr_detr ? Implicit() : Explicit(),
         sgs_nh_pressure_mode = implicit_sgs_nh_pressure ? Implicit() :
                                Explicit(),
+        sgs_vertdiff_mode = implicit_sgs_vertdiff ? Implicit() : Explicit(),
         sgs_mf_mode = implicit_sgs_mass_flux ? Implicit() : Explicit(),
         smagorinsky_lilly = get_smagorinsky_lilly_model(parsed_args),
 
@@ -482,6 +486,7 @@ function get_jacobian(ode_algo, Y, atmos, parsed_args)
             DerivativeFlag(atmos.sgs_entr_detr_mode),
             DerivativeFlag(atmos.sgs_mf_mode),
             DerivativeFlag(atmos.sgs_nh_pressure_mode),
+            DerivativeFlag(atmos.sgs_vertdiff_mode),
             parsed_args["approximate_linear_solve_iters"],
         )
         parsed_args["use_auto_jacobian"] ?
diff --git a/src/solver/types.jl b/src/solver/types.jl
@@ -575,12 +575,13 @@ end
 
 Groups turbulence convection-related models and types.
 """
-Base.@kwdef struct AtmosTurbconv{EDMFX, TCM, SAM, SEDM, SNPM, SMM, SL}
+Base.@kwdef struct AtmosTurbconv{EDMFX, TCM, SAM, SEDM, SNPM, SVM, SMM, SL}
     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
 end
@@ -807,7 +808,7 @@ Internal testing and calibration components for single-column setups:
 ## AtmosTurbconv
 - `edmfx_model`: EDMFXModel()
 - `turbconv_model`: nothing, PrognosticEDMFX(), DiagnosticEDMFX(), EDOnlyEDMFX()
-- `sgs_adv_mode`, `sgs_entr_detr_mode`, `sgs_nh_pressure_mode`, `sgs_mf_mode`: Explicit(), Implicit()
+- `sgs_adv_mode`, `sgs_entr_detr_mode`, `sgs_nh_pressure_mode`, `sgs_vertdiff_mode`, `sgs_mf_mode`: Explicit(), Implicit()
 - `smagorinsky_lilly`: nothing or SmagorinskyLilly()
 
 ## AtmosGravityWave
