updating linear solver

Author: Olivia Alcabes

src/prognostic_equations/implicit/approx_jacobian.jl

@@ -257,67 +257,61 @@ function jacobian_cache(alg::ApproxJacobian, Y, atmos)
         sgs_massflux_blocks...,
     )
 
-    names₁_group₁ = (@name(c.ρ), sfc_if_available...)
-    names₁_group₂ = (available_tracer_names..., ρatke_if_available...)
-    names₁_group₃ = (@name(c.ρe_tot),)
-    names₁ = (
-        names₁_group₁...,
-        names₁_group₂...,
-        names₁_group₃...,
+    mass_and_surface_names = (@name(c.ρ), sfc_if_available...)
+    available_scalar_names = (
+        mass_and_surface_names...,
+        ρq_tot_if_available...,
+        @name(c.ρe_tot),
+        available_condensate_names...,
+        ρatke_if_available...,
         available_sgs_scalar_names...,
     )
 
-    alg₂ = MatrixFields.BlockLowerTriangularSolve(
+    velocity_alg = MatrixFields.BlockLowerTriangularSolve(
         @name(c.uₕ),
         sgs_u³_if_available...,
     )
-    solver_alg =
+    full_alg =
         if use_derivative(diffusion_flag) ||
            use_derivative(sgs_advection_flag) ||
            !(atmos.moisture_model isa DryModel)
-            alg₁_subalg₂ =
+            gs_scalar_subalg = if !(atmos.moisture_model isa DryModel)
+                MatrixFields.BlockLowerTriangularSolve(@name(c.ρq_tot))
+            else
+                MatrixFields.BlockDiagonalSolve()
+            end
+            scalar_subalg =
                 if atmos.turbconv_model isa PrognosticEDMFX &&
                    use_derivative(sgs_advection_flag)
-                    diff_subalg =
-                        use_derivative(diffusion_flag) ?
-                        (;
-                            alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                                names₁_group₂...,
-                            )
-                        ) : (;)
-                    (;
+                    MatrixFields.BlockLowerTriangularSolve(
+                        available_sgs_tracer_names...;
                         alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                            available_sgs_tracer_names...;
+                            @name(c.sgsʲs.:(1).mse);
                             alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                                @name(c.sgsʲs.:(1).mse);
-                                alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                                    @name(c.sgsʲs.:(1).ρa);
-                                    diff_subalg...,
-                                ),
+                                @name(c.sgsʲs.:(1).ρa);
+                                alg₂ = gs_scalar_subalg,
                             ),
-                        )
+                        ),
                     )
                 else
-                    is_in_Y(@name(c.ρq_tot)) ?
-                    (;
-                        alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                            names₁_group₂...,
-                        )
-                    ) : (;)
+                    gs_scalar_subalg
                 end
-            alg₁ = MatrixFields.BlockLowerTriangularSolve(
-                names₁_group₁...;
-                alg₁_subalg₂...,
+            scalar_alg = MatrixFields.BlockLowerTriangularSolve(
+                mass_and_surface_names...;
+                alg₂ = scalar_subalg,
             )
             MatrixFields.ApproximateBlockArrowheadIterativeSolve(
-                names₁...;
-                alg₁,
-                alg₂,
+                available_scalar_names...;
+                alg₁ = scalar_alg,
+                alg₂ = velocity_alg,
                 P_alg₁ = MatrixFields.MainDiagonalPreconditioner(),
                 n_iters = alg.approximate_solve_iters,
             )
         else
-            MatrixFields.BlockArrowheadSolve(names₁...; alg₂)
+            MatrixFields.BlockArrowheadSolve(
+                available_scalar_names...;
+                alg₂ = velocity_alg,
+            )
         end
 
     temp_matrix = (matrix .+ identity_matrix(matrix, Y)) ./ FT(1)