ASMStar: enable patch coloring

demos/fast_diagonalisation/fast_diagonalisation_poisson.py.rst

@@ -95,14 +95,17 @@ using a sparse direct LU factorization. ::
 .. note::
     On this Cartesian mesh, the sparse operator constructed by :class:`~.FDMPC`
     corresponds to the original operator. This is no longer the case with non-Cartesian
-    meshes or more general PDEs, as the FDM basis will fail to diagonalise the
-    problem.  For such cases, :class:`~.FDMPC` will produce a sparse
+    meshes or more general PDEs, as the FDM basis only diagonalises very specific
+    problems.  For such cases, :class:`~.FDMPC` will produce a sparse
     approximation of the original operator.
 
 Moving on to a more complicated solver, we'll employ a two-level solver with
 the lowest-order coarse space via :class:`~.P1PC`.  As the fine level
 relaxation we define an additive Schwarz method on vertex-star patches
-implemented via :class:`~.ASMExtrudedStarPC` as we have an extruded mesh::
+implemented via :class:`~.ASMExtrudedStarPC` as we have an extruded mesh.
+In addition we specify `"use_coloring"` to group non-overlapping subsets of
+patches into sparse block-diagonal matrices via a mesh coloring, which reduces
+the overhead of calling many KSP solves for each patch.::
 
   asm_params = {
       "pc_type": "python",
@@ -114,7 +117,8 @@ implemented via :class:`~.ASMExtrudedStarPC` as we have an extruded mesh::
           "ksp_type": "chebyshev",
           "pc_type": "python",
           "pc_python_type": "firedrake.ASMExtrudedStarPC",
-          "sub_sub_pc_type": "lu",
+          "pc_star_use_coloring": True,
+          "pc_star_sub_sub_pc_type": "lu",
       },
   }
 

demos/patch/poisson_mg_patches.py.rst

@@ -116,16 +116,27 @@ in dense format. ::
 :class:`~.ASMStarPC`, on the other hand, does no re-discretization, but extracts the
 submatrices for each patch from the already-assembled global stiffness matrix.
 
+The `"tinyasm"` backend uses LAPACK to invert all the patch operators, which is ideal
+for low-order discretizations ::
 
-The tinyasm backend uses LAPACK to invert all the patch operators.  If this option
-is not specified, PETSc's ASM framework will set up a small KSP for each patch.
+  tinyasm_relax = mg_params({
+      "pc_type": "python",
+      "pc_python_type": "firedrake.ASMStarPC",
+      "pc_star_backend": "tinyasm"})
+
+If the backend option is not specified, PETSc's ASM framework will set up a KSP for each patch.
 This can be useful when the patches become larger and one wants to use a sparse
-direct solver or a Krylov iteration on each one. ::
+direct solver or a Krylov iteration on each one.
+
+Moreover, the option `"use_coloring"` applies a mesh coloring to combine
+subsets of non-overlapping patches into sparse block-diagonal matrices. This
+results in a mathematically equivalent preconditioner, while reducing the
+overhead costs of calling the sparse factorization library many times. ::
 
-  asm_relax = mg_params({
+  color_relax = mg_params({
       "pc_type": "python",
       "pc_python_type": "firedrake.ASMStarPC",
-      "pc_star_backend": "tinyasm"})
+      "pc_star_use_coloring": True})
 
 Now, for each parameter choice, we report the iteration count for the Poisson problem
 over a range of polynomial degrees.  We see that the Jacobi relaxation leads to growth
@@ -134,7 +145,8 @@ latter options do the same operations, just via different code paths. ::
 
   names = {"Jacobi": jacobi_relax,
            "Patch": patch_relax,
-           "ASM Star": asm_relax}
+           "Tiny ASM": tinyasm_relax,
+           "Color ASM": color_relax}
 
   for name, method in names.items():
       print(name)
@@ -158,7 +170,7 @@ For Jacobi, we expect output such as
    7         19
 ======== ================
 
-While for either :class:`~.PatchPC` or :class:`~.ASMStarPC`, we expect
+While for either :class:`~.PatchPC` or :class:`~.ASMStarPC` (with dense inversion or with coloring + sparse factorization), we expect
 
 ======== ================
  Degree    Iterations