making implicit solve much more readable

Author: Olivia Alcabes

src/prognostic_equations/implicit/implicit_solver.jl

@@ -184,14 +184,30 @@ function ImplicitEquationJacobian(
         is_in_Y(@name(c.sgs⁰.ρatke)) ? (@name(c.sgs⁰.ρatke),) : ()
     sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : ()
 
-    tracer_names = (
-        @name(c.ρq_tot),
-        @name(c.ρq_liq),
-        @name(c.ρq_ice),
-        @name(c.ρq_rai),
-        @name(c.ρq_sno),
+    condensate_names =
+        (@name(c.ρq_liq), @name(c.ρq_ice), @name(c.ρq_rai), @name(c.ρq_sno))
+    available_condensate_names =
+        MatrixFields.unrolled_filter(is_in_Y, condensate_names)
+    available_tracer_names =
+        (ρq_tot_if_available..., available_condensate_names...)
+
+    sgs_tracer_names = (
+        @name(c.sgsʲs.:(1).q_tot),
+        @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),
     )
-    available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names)
+    available_sgs_tracer_names =
+        MatrixFields.unrolled_filter(is_in_Y, sgs_tracer_names)
+
+    sgs_scalar_names =
+        (sgs_tracer_names..., @name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).ρa))
+    available_sgs_scalar_names =
+        MatrixFields.unrolled_filter(is_in_Y, sgs_scalar_names)
+
+    sgs_u³_if_available =
+        is_in_Y(@name(f.sgsʲs.:(1).u₃)) ? (@name(f.sgsʲs.:(1).u₃),) : ()
 
     # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0,
     # which means that multiplying inv(-1) by a Float32 will yield a Float64.
@@ -266,21 +282,8 @@ function ImplicitEquationJacobian(
         )
     end
 
-    sgs_scalar_names = (
-        @name(c.sgsʲs.:(1).q_tot),
-        @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),
-        @name(c.sgsʲs.:(1).mse),
-        @name(c.sgsʲs.:(1).ρa)
-    )
-    available_sgs_scalar_names =
-        MatrixFields.unrolled_filter(is_in_Y, sgs_scalar_names)
-
     sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX
         @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1
-
         if use_derivative(sgs_advection_flag)
             (
                 MatrixFields.unrolled_map(
@@ -353,79 +356,65 @@ function ImplicitEquationJacobian(
         sgs_massflux_blocks...,
     )
 
-    sgs_u³_names_if_available = if atmos.turbconv_model isa PrognosticEDMFX
-        (@name(f.sgsʲs.:(1).u₃),)
-    else
-        ()
-    end
-
-    names₁_group₁ = (@name(c.ρ), sfc_if_available...)
-    names₁_group₂ = (available_tracer_names..., ρatke_if_available...)
-    names₁_group₃ = (@name(c.ρe_tot),)
-    names₁ = (
-        names₁_group₁...,
+    mass_and_surface_names = (@name(c.ρ), sfc_if_available...)
+    available_scalar_names = (
+        mass_and_surface_names...,
+        available_tracer_names...,
+        @name(c.ρe_tot),
+        ρatke_if_available...,
         available_sgs_scalar_names...,
-        names₁_group₂...,
-        names₁_group₃...,
     )
 
-    alg₂ = MatrixFields.BlockLowerTriangularSolve(
+    velocity_alg = MatrixFields.BlockLowerTriangularSolve(
         @name(c.uₕ),
-        sgs_u³_names_if_available...,
+        sgs_u³_if_available...,
     )
-    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(
-                            # TODO: What needs to be changed here for 1M?
-                            @name(c.sgsʲs.:(1).q_tot);
+                            @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 = approximate_solve_iters,
             )
         else
-            MatrixFields.BlockArrowheadSolve(names₁...; alg₂)
+            MatrixFields.BlockArrowheadSolve(
+                available_scalar_names...;
+                alg₂ = velocity_alg,
+            )
         end
 
     return ImplicitEquationJacobian(
         matrix,
-        MatrixFields.FieldMatrixSolver(alg, matrix, Y),
+        MatrixFields.FieldMatrixSolver(full_alg, matrix, Y),
         diffusion_flag,
         topography_flag,
         sgs_advection_flag,