From c90bb4628bf41124af453b23998623bf2692a9e1 Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Fri, 29 Mar 2024 14:48:55 -0700
Subject: [PATCH] fix

---
 .../implicit/implicit_solver.jl               | 20 +++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl
index 1e65bcc3574..1847ded9889 100644
--- a/src/prognostic_equations/implicit/implicit_solver.jl
+++ b/src/prognostic_equations/implicit/implicit_solver.jl
@@ -9,7 +9,10 @@ use_derivative(::UseDerivative) = true
 use_derivative(::IgnoreDerivative) = false
 
 """
-    ImplicitEquationJacobian(Y, atmos; approximate_solve_iters, diffusion_flag, topography_flag, sgs_advection_flag, transform_flag)
+    ImplicitEquationJacobian(
+        Y, atmos; 
+        approximate_solve_iters, diffusion_flag, topography_flag, sgs_advection_flag, transform_flag
+    )
 
 A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the
 implicit step with the state ``Y``.
@@ -76,7 +79,7 @@ is reached.
 - `diffusion_flag::DerivativeFlag`: whether the derivative of the
   diffusion tendency with respect to the quantities being diffused should be
   computed or approximated as 0; must be either `UseDerivative()` or
-  `Ignoreerivative()` instead of a `Bool` to ensure type-stability
+  `IgnoreDerivative()` instead of a `Bool` to ensure type-stability
 - `topography_flag::DerivativeFlag`: whether the derivative of vertical
   contravariant velocity with respect to horizontal covariant velocity should be
   computed or approximated as 0; must be either `UseDerivative()` or
@@ -84,7 +87,7 @@ is reached.
 - `sgs_advection_flag::DerivativeFlag`: whether the derivative of the
   subgrid-scale advection tendency with respect to the updraft quantities should be
   computed or approximated as 0; must be either `UseDerivative()` or
-  `Ignoreerivative()` instead of a `Bool` to ensure type-stability
+  `IgnoreDerivative()` instead of a `Bool` to ensure type-stability
 - `transform_flag::Bool`: whether the error should be transformed from ``E(Y)``
   to ``E(Y)/Δt``, which is required for non-Rosenbrock timestepping schemes from
   OrdinaryDiffEq.jl
@@ -300,6 +303,13 @@ function ImplicitEquationJacobian(
             alg₁_subalg₂ =
                 if atmos.turbconv_model isa PrognosticEDMFX &&
                    use_derivative(sgs_advection_flag)
+                    diff_subalg =
+                        use_derivative(diffusion_flag) ?
+                        (;
+                            alg₂ = MatrixFields.BlockLowerTriangularSolve(
+                                names₁_group₂...,
+                            )
+                        ) : (;)
                     (;
                         alg₂ = MatrixFields.BlockLowerTriangularSolve(
                             @name(c.sgsʲs.:(1).q_tot);
@@ -308,9 +318,7 @@ function ImplicitEquationJacobian(
                                 alg₂ = MatrixFields.BlockLowerTriangularSolve(
                                     @name(c.sgsʲs.:(1).ρa),
                                     @name(f.sgsʲs.:(1).u₃);
-                                    alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                                        names₁_group₂...,
-                                    ),
+                                    diff_subalg...,
                                 ),
                             ),
                         )