Merge pull request #3909 from CliMA/zs/updraft_diffusion

szy21 · web-flow · commit 57f50f2c34bc · 2025-09-12T22:06:13.000Z
add vertical diffusion of updraft
diff --git a/calibration/experiments/gcm_driven_scm/model_config_diagnostic.yml b/calibration/experiments/gcm_driven_scm/model_config_diagnostic.yml
@@ -14,7 +14,6 @@ edmfx_detr_model: "PiGroups"
 edmfx_nh_pressure: true
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
-edmfx_nh_pressure: true
 moist: "equil"
 call_cloud_diagnostics_per_stage: true
 config: "column"
diff --git a/calibration/experiments/gcm_driven_scm/model_config_prognostic.yml b/calibration/experiments/gcm_driven_scm/model_config_prognostic.yml
@@ -15,6 +15,7 @@ edmfx_detr_model: "PiGroups"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: false
 prognostic_tke: true
 moist: "equil"
diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
@@ -366,6 +366,9 @@ edmfx_nh_pressure:
   help: "If set to true, it switches on EDMFX pressure drag closure.  For prognostic EDMF, this only controls the drag term 
   in the pressure closure. The buoyancy term is always applied. [`true`, `false` (default)]"
   value: false
+edmfx_vertical_diffusion:
+  help: "If set to true, it switches on vertical diffusion of prognostic EDMFX updrafts. [`true`, `false` (default)]"
+  value: false
 edmfx_entr_model:
   help: "EDMFX entrainment closure.  [`nothing` (default), `PiGroups`, `Generalized`]"
   value: ~
diff --git a/config/model_configs/prognostic_edmfx_aquaplanet.yml b/config/model_configs/prognostic_edmfx_aquaplanet.yml
@@ -12,6 +12,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_dense_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_dense_autodiff.yml
@@ -14,6 +14,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_gpu.yml b/config/model_configs/prognostic_edmfx_aquaplanet_gpu.yml
@@ -14,6 +14,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_gpu_dense_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_gpu_dense_autodiff.yml
@@ -14,6 +14,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_gpu_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_gpu_sparse_autodiff.yml
@@ -14,6 +14,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_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_aquaplanet_sparse_autodiff.yml
@@ -14,6 +14,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_bomex_box.yml b/config/model_configs/prognostic_edmfx_bomex_box.yml
@@ -14,6 +14,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_bomex_column.yml b/config/model_configs/prognostic_edmfx_bomex_column.yml
@@ -13,6 +13,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_bomex_column_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_bomex_column_sparse_autodiff.yml
@@ -20,6 +20,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 precip_model: 0M
diff --git a/config/model_configs/prognostic_edmfx_bomex_fixtke_column.yml b/config/model_configs/prognostic_edmfx_bomex_fixtke_column.yml
@@ -10,6 +10,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: false
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml b/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml
@@ -17,6 +17,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 precip_model: 0M
diff --git a/config/model_configs/prognostic_edmfx_bomex_pigroup_column.yml b/config/model_configs/prognostic_edmfx_bomex_pigroup_column.yml
@@ -10,6 +10,7 @@ edmfx_detr_model: "PiGroups"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 cloud_model: "grid_scale"
diff --git a/config/model_configs/prognostic_edmfx_bomex_stretched_column.yml b/config/model_configs/prognostic_edmfx_bomex_stretched_column.yml
@@ -13,6 +13,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_diurnal_scm_imp.yml b/config/model_configs/prognostic_edmfx_diurnal_scm_imp.yml
@@ -26,6 +26,7 @@ edmfx_detr_model: "PiGroups"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 precip_model: "0M"
diff --git a/config/model_configs/prognostic_edmfx_dycoms_rf01_column.yml b/config/model_configs/prognostic_edmfx_dycoms_rf01_column.yml
@@ -13,6 +13,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: equil
diff --git a/config/model_configs/prognostic_edmfx_dycoms_rf02_column.yml b/config/model_configs/prognostic_edmfx_dycoms_rf02_column.yml
@@ -12,6 +12,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "nonequil"
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
@@ -20,6 +20,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "nonequil"
diff --git a/config/model_configs/prognostic_edmfx_gabls_column.yml b/config/model_configs/prognostic_edmfx_gabls_column.yml
@@ -11,6 +11,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_gabls_column_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_gabls_column_sparse_autodiff.yml
@@ -18,6 +18,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_gcmdriven_column.yml b/config/model_configs/prognostic_edmfx_gcmdriven_column.yml
@@ -14,6 +14,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_rico_column.yml b/config/model_configs/prognostic_edmfx_rico_column.yml
@@ -12,6 +12,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "nonequil"
diff --git a/config/model_configs/prognostic_edmfx_rico_column_2M.yml b/config/model_configs/prognostic_edmfx_rico_column_2M.yml
@@ -12,6 +12,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "nonequil"
diff --git a/config/model_configs/prognostic_edmfx_rico_implicit_column.yml b/config/model_configs/prognostic_edmfx_rico_implicit_column.yml
@@ -15,6 +15,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "nonequil"
diff --git a/config/model_configs/prognostic_edmfx_simpleplume_column.yml b/config/model_configs/prognostic_edmfx_simpleplume_column.yml
@@ -12,6 +12,7 @@ edmfx_detr_model: "SmoothArea"
 edmfx_sgs_mass_flux: false
 edmfx_sgs_diffusive_flux: false
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: false
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_soares_column.yml b/config/model_configs/prognostic_edmfx_soares_column.yml
@@ -12,6 +12,7 @@ edmfx_detr_model: "SmoothArea"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "equil"
diff --git a/config/model_configs/prognostic_edmfx_trmm_column.yml b/config/model_configs/prognostic_edmfx_trmm_column.yml
@@ -12,6 +12,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: nonequil
diff --git a/config/model_configs/prognostic_edmfx_trmm_column_0M.yml b/config/model_configs/prognostic_edmfx_trmm_column_0M.yml
@@ -15,6 +15,7 @@ edmfx_detr_model: "SmoothArea"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: equil
diff --git a/config/model_configs/prognostic_edmfx_trmm_column_sparse_autodiff.yml b/config/model_configs/prognostic_edmfx_trmm_column_sparse_autodiff.yml
@@ -18,6 +18,7 @@ edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 moist: nonequil
diff --git a/config/model_configs/prognostic_edmfx_tv_era5driven_column.yml b/config/model_configs/prognostic_edmfx_tv_era5driven_column.yml
@@ -19,6 +19,7 @@ edmfx_detr_model: "PiGroups"
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 edmfx_nh_pressure: true
+edmfx_vertical_diffusion: true
 edmfx_filter: true
 prognostic_tke: true
 precip_model: "0M"
diff --git a/config/perf_configs/bm_aquaplanet_progedmf.yml b/config/perf_configs/bm_aquaplanet_progedmf.yml
@@ -20,6 +20,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/perf_configs/bm_aquaplanet_progedmf_dense_autodiff.yml b/config/perf_configs/bm_aquaplanet_progedmf_dense_autodiff.yml
@@ -23,6 +23,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/perf_configs/bm_aquaplanet_progedmf_sparse_autodiff.yml b/config/perf_configs/bm_aquaplanet_progedmf_sparse_autodiff.yml
@@ -24,6 +24,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/reproducibility_tests/ref_counter.jl b/reproducibility_tests/ref_counter.jl
@@ -1,4 +1,4 @@
-261
+262
 
 # **README**
 #
@@ -20,6 +20,9 @@
 
 
 #=
+262
+- Add vertical diffusion to prognostic EDMF updrafts
+
 261
 - Ignore the implicit solver Jacobian for the tracer-density block from diffusion
 
diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl
@@ -399,9 +399,10 @@ function edmfx_sgs_diffusive_flux_tendency!(
 )
     FT = Spaces.undertype(axes(Y.c))
     (; dt, params) = p
+    thermo_params = CAP.thermodynamics_params(params)
     turbconv_params = CAP.turbconv_params(params)
     c_d = CAP.tke_diss_coeff(turbconv_params)
-    (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
+    (; ᶜK, ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜts) = p.precomputed
     (; ρatke_flux) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
     ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
@@ -425,18 +426,24 @@ function edmfx_sgs_diffusive_flux_tendency!(
         )
         ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec))
         ᶠρaK_h = p.scratch.ᶠtemp_scalar
-        @. ᶠρaK_h = ᶠinterp(ᶜρa⁰) * ᶠinterp(ᶜK_h)
+        @. ᶠρaK_h = ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h)
         ᶠρaK_u = p.scratch.ᶠtemp_scalar_2
-        @. ᶠρaK_u = ᶠinterp(ᶜρa⁰) * ᶠinterp(ᶜK_u)
+        @. ᶠρaK_u = ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_u)
 
         # Total enthalpy diffusion
         ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(
             top = Operators.SetValue(C3(FT(0))),
             bottom = Operators.SetValue(C3(FT(0))),
         )
 
-        ᶜmse⁰ = ᶜspecific_env_mse(Y, p)
-        @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜmse⁰ + ᶜK⁰)))
+        ᶜh_tot = @. lazy(
+            TD.total_specific_enthalpy(
+                thermo_params,
+                ᶜts,
+                specific(Y.c.ρe_tot, Y.c.ρ),
+            ),
+        )
+        @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜh_tot)))
         if use_prognostic_tke(turbconv_model)
             # Turbulent TKE transport (diffusion)
             ᶜdivᵥ_ρatke = Operators.DivergenceF2C(
@@ -464,8 +471,8 @@ function edmfx_sgs_diffusive_flux_tendency!(
                 top = Operators.SetValue(C3(FT(0))),
                 bottom = Operators.SetValue(C3(FT(0))),
             )
-            ᶜq_tot⁰ = ᶜspecific_env_value(@name(q_tot), Y, p)
-            @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜq_tot⁰)))
+            @. ᶜρχₜ_diffusion =
+                ᶜ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
         end
@@ -489,18 +496,16 @@ function edmfx_sgs_diffusive_flux_tendency!(
         )
         MatrixFields.unrolled_foreach(cloud_tracers) do (ρχ_name, χ_name)
             MatrixFields.has_field(Y, ρχ_name) || return
-            ᶜχ⁰ = ᶜspecific_env_value(χ_name, Y, p)
-            @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq(-(ᶠρaK_h * ᶠgradᵥ(ᶜχ⁰)))
+            ᶜρχ = MatrixFields.get_field(Y, ρχ_name)
+            @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq(-(ᶠρaK_h * ᶠgradᵥ(specific(ᶜρχ, Y.c.ρ))))
             ᶜρχₜ = MatrixFields.get_field(Yₜ, ρχ_name)
             @. ᶜρχₜ -= ᶜρχₜ_diffusion
         end
-        # TODO - do I need to change anything in the implicit solver
-        # to include the α_vert_diff_tracer?
         MatrixFields.unrolled_foreach(precip_tracers) do (ρχ_name, χ_name)
             MatrixFields.has_field(Y, ρχ_name) || return
-            ᶜχ⁰ = ᶜspecific_env_value(χ_name, Y, p)
+            ᶜρχ = MatrixFields.get_field(Y, ρχ_name)
             @. ᶜρχₜ_diffusion =
-                ᶜdivᵥ_ρq(-(ᶠρaK_h * α_vert_diff_tracer * ᶠgradᵥ(ᶜχ⁰)))
+                ᶜdivᵥ_ρq(-(ᶠρaK_h * α_vert_diff_tracer * ᶠgradᵥ(specific(ᶜρχ, Y.c.ρ))))
             ᶜρχₜ = MatrixFields.get_field(Yₜ, ρχ_name)
             @. ᶜρχₜ -= ᶜρχₜ_diffusion
         end
diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl
@@ -640,8 +640,6 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
 
         MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, _)
             MatrixFields.has_field(Y, ρχ_name) || return
-            ᶜρχ = MatrixFields.get_field(Y, ρχ_name)
-            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
             ∂ᶜρχ_err_∂ᶜρ = matrix[ρχ_name, @name(c.ρ)]
             ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name]
             ᶜtridiagonal_matrix_scalar = ifelse(
@@ -690,13 +688,13 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             @. ∂ᶜρatke⁰_err_∂ᶜρ =
                 dtγ * (
                     DiagonalMatrixRow(ᶜdissipation_matrix_diagonal)
-                ) ⋅ DiagonalMatrixRow(ᶜtke⁰ / ᶜρa⁰)
+                ) ⋅ DiagonalMatrixRow(ᶜtke⁰ / Y.c.ρ)
             @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰ =
                 dtγ * (
                     (
                         ᶜdiffusion_u_matrix -
                         DiagonalMatrixRow(ᶜdissipation_matrix_diagonal)
-                    ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰) - DiagonalMatrixRow(
+                    ) ⋅ DiagonalMatrixRow(1 / Y.c.ρ) - DiagonalMatrixRow(
                         tke_dissipation_rate_tendency(
                             ᶜtke⁰,
                             ᶜmixing_length_field,
diff --git a/src/prognostic_equations/mass_flux_closures.jl b/src/prognostic_equations/mass_flux_closures.jl
diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
diff --git a/src/solver/types.jl b/src/solver/types.jl
