Fast diagonalization demo

demos/demo_references.bib

@@ -308,3 +308,25 @@ @article{Ricker:1953
   publisher={Society of Exploration Geophysicists},
   doi={https://doi.org/10.1190/1.1437843}
 }
+
+@article{Brubeck2022,
+  author        = {Brubeck, Pablo D. and Farrell, Patrick E.},
+  doi           = {10.1137/21M1444187},
+  journal       = {SIAM J. Sci. Comput.},
+  number        = {5},
+  pages         = {A2991-A3017},
+  title         = {A scalable and robust vertex-star relaxation for high-order {FEM}},
+  volume        = {44},
+  year          = {2022}
+}
+
+@article{Brubeck2024,
+  author        = {Brubeck, Pablo D. and Farrell, Patrick E.},
+  doi           = {10.1137/22M1537370},
+  journal       = {SIAM J. Sci. Comput.},
+  number        = {3},
+  pages         = {A1549-A1573},
+  title         = {{Multigrid solvers for the de Rham complex with optimal complexity in polynomial degree}},
+  volume        = {46},
+  year          = {2024}
+}

demos/fast_diagonalisation/fast_diagonalisation_poisson.py.rst

@@ -0,0 +1,200 @@
+Using fast diagonalisation solvers in Firedrake
+===============================================
+
+Contributed by `Pablo Brubeck <https://www.maths.ox.ac.uk/people/pablo.brubeckmartinez/>`_.
+
+In this demo we show how to efficiently solve the Poisson equation using
+high-order tensor-product elements. This is accomplished through a special
+basis, obtained from the fast diagonalisation method (FDM). We first construct
+an auxiliary operator that is sparse in this basis, with as many zeros as a
+low-order method.  We then combine this with an additive Schwarz method.
+Finally, we show how to do static condensation using fieldsplit. A detailed
+description of these solvers is found in :cite:`Brubeck2022` and
+:cite:`Brubeck2024`.
+
+
+Creating an extruded mesh
+-------------------------
+
+The fast diagonalisation method produces a basis of discrete eigenfunctions.
+These are polynomials, and can be efficiently computed on tensor
+product-elements by solving an eigenproblem on the interval. Therefore, we will
+require quadrilateral or hexahedral meshes.  Currently, the solver only supports
+extruded hexahedral meshes, so we must create an :func:`~.ExtrudedMesh`. ::
+
+  from firedrake import *
+
+  base = UnitSquareMesh(8, 8, quadrilateral=True)
+  mesh = ExtrudedMesh(base, 8)
+
+
+Defining the problem: the Poisson equation
+------------------------------------------
+
+Having defined the mesh we now need to set up our problem.  The crucial step
+for fast diagonalisation is a special choice of basis functions. We obtain them
+by passing ``variant="fdm"`` to the :func:`~.FunctionSpace` constructor.  The
+solvers in this demo work also with other element variants, but each iteration
+would involve an additional a basis transformation.  To stress-test the solver,
+we prescribe a random :class:`~.Cofunction` as right-hand side.
+
+We'll demonstrate a few different sets of solver parameters, so let's define a
+function that takes in set of parameters and uses them on a
+:class:`~.LinearVariationalSolver`. ::
+
+
+  def run_solve(degree, parameters):
+      V = FunctionSpace(mesh, "Q", degree, variant="fdm")
+      u = TrialFunction(V)
+      v = TestFunction(V)
+      uh = Function(V)
+      a = inner(grad(u), grad(v)) * dx
+      bcs = [DirichletBC(V, 0, sub) for sub in ("on_boundary", "top", "bottom")]
+
+      rg = RandomGenerator(PCG64(seed=123456789))
+      L = rg.uniform(V.dual(), -1, 1)
+
+      problem = LinearVariationalProblem(a, L, uh, bcs=bcs)
+      solver = LinearVariationalSolver(problem, solver_parameters=parameters)
+      solver.solve()
+      return solver.snes.getLinearSolveIterations()
+
+Specifying the solver
+---------------------
+
+The solver avoids the assembly of a matrix with dense element submatrices, and
+instead applies a matrix-free conjugate gradient method with a preconditioner
+obtained by assembling a sparse matrix.  This is done through the python type
+preconditioner :class:`~.FDMPC`.  We define a function that enables us to
+compose :class:`~.FDMPC` with an inner relaxation. ::
+
+
+  def fdm_params(relax):
+      return {
+          "mat_type": "matfree",
+          "ksp_type": "cg",
+          "pc_type": "python",
+          "pc_python_type": "firedrake.FDMPC",
+          "fdm": relax,
+      }
+
+Let's start with our first test.  We'll confirm a working solve by
+using a sparse direct LU factorization. ::
+
+
+  lu_params = {
+      "pc_type": "lu",
+      "pc_factor_mat_solver_type": "mumps",
+  }
+
+  fdm_lu_params = fdm_params(lu_params)
+  its = run_solve(5, fdm_lu_params)
+  print(f"LU iterations {its}")
+
+
+.. 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
+    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::
+
+  asm_params = {
+      "pc_type": "python",
+      "pc_python_type": "firedrake.P1PC",
+      "pmg_mg_coarse_mat_type": "aij",
+      "pmg_mg_coarse": lu_params,
+      "pmg_mg_levels": {
+          "ksp_max_it": 1,
+          "ksp_type": "chebyshev",
+          "pc_type": "python",
+          "pc_python_type": "firedrake.ASMExtrudedStarPC",
+          "sub_sub_pc_type": "lu",
+      },
+  }
+
+  fdm_asm_params = fdm_params(asm_params)
+
+  print("FDM + ASM")
+  print("Degree\tIterations")
+  for degree in range(3, 6):
+      its = run_solve(degree, fdm_asm_params)
+      print(f"{degree}\t{its}")
+
+We observe degree-independent iteration counts:
+
+======== ============
+ Degree   Iterations
+======== ============
+  3        13
+  4        13
+  5        12
+======== ============
+
+Static condensation
+-------------------
+
+Finally, we construct :class:`~.FDMPC` solver parameters using static
+condensation.  The fast diagonalisation basis diagonalises the operator on cell
+interiors. So we define a solver that splits the interior and facet degrees of
+freedom via :class:`~.FacetSplitPC` and fieldsplit options.  We set the option
+``fdm_static_condensation`` to tell :class:`~.FDMPC` to assemble a 2-by-2 block
+preconditioner where the lower-right block is replaced by the Schur complement
+resulting from eliminating the interior degrees of freedom.  The Krylov
+solver is posed on the full set of degrees of freedom, and the preconditioner
+applies a symmetrized multiplicative sweep on the interior and the facet
+degrees of freedom. In general, we are not able to fully eliminate the
+interior, as the sparse operator constructed by :class:`~.FDMPC` is only an
+approximation on non-Cartesian meshes.  We apply point-Jacobi on the interior
+block, and the two-level additive Schwarz method on the facets. ::
+
+
+  def fdm_static_condensation_params(relax):
+      return {
+          "mat_type": "matfree",
+          "ksp_type": "cg",
+          "pc_type": "python",
+          "pc_python_type": "firedrake.FacetSplitPC",
+          "facet_pc_type": "python",
+          "facet_pc_python_type": "firedrake.FDMPC",
+          "facet_fdm_static_condensation": True,
+          "facet_fdm_pc_use_amat": False,
+          "facet_fdm_pc_type": "fieldsplit",
+          "facet_fdm_pc_fieldsplit_type": "symmetric_multiplicative",
+          "facet_fdm_fieldsplit_ksp_type": "preonly",
+          "facet_fdm_fieldsplit_0_pc_type": "jacobi",
+          "facet_fdm_fieldsplit_1": relax,
+      }
+
+
+  fdm_sc_asm_params = fdm_static_condensation_params(asm_params)
+
+  print('FDM + SC + ASM')
+  print("Degree\tIterations")
+  for degree in range(3, 6):
+      its = run_solve(degree, fdm_sc_asm_params)
+      print(f"{degree}\t{its}")
+
+We also observe degree-independent iteration counts:
+
+======== ============
+ Degree   Iterations
+======== ============
+  3        10
+  4        10
+  5        10
+======== ============
+
+A runnable python version of this demo can be found :demo:`here
+<fast_diagonalisation_poisson.py>`.
+
+
+.. rubric:: References
+
+.. bibliography:: demo_references.bib
+   :filter: docname in docnames

docs/source/_static/firedrake-apps.bib

@@ -198,17 +198,6 @@ @article{Braun2017
   year          = {2017}
 }
 
-@article{Brubeck2022,
-  author        = {Brubeck, Pablo D. and Farrell, Patrick E.},
-  doi           = {10.1137/21M1444187},
-  journal       = {SIAM J. Sci. Comput.},
-  number        = {5},
-  pages         = {A2991-A3017},
-  title         = {A scalable and robust vertex-star relaxation for high-order {FEM}},
-  volume        = {44},
-  year          = {2022}
-}
-
 @misc{Brugnoli2020,
   author        = {Brugnoli, Andrea and Alazar, Daniel and Pommier-Budinger, Valérie and Matignon, Denis},
   eprint        = {2002.12816},

docs/source/advanced_tut.rst

@@ -30,3 +30,4 @@ element systems.
    Vertex/edge star multigrid relaxation for H(div).<demos/hdiv_riesz_star.py>
    Auxiliary space patch relaxation multigrid for H(curl).<demos/hcurl_riesz_star.py>
    Steady Boussinesq problem with integral constraints.<demos/boussinesq.py>
+   Preconditioning using fast diagonalisation.<demos/fast_diagonalisation_poisson.py>

docs/source/preconditioning.rst

@@ -138,7 +138,7 @@ operator instead.
    implemented for quadrilateral and hexahedral cells. The assembled
    matrix becomes as sparse as a low-order refined preconditioner, to
    which one may apply other preconditioners such as :class:`.ASMStarPC` or
-   :class:`.ASMExtrudedStarPC`. See details in :cite:`Brubeck2022`.
+   :class:`.ASMExtrudedStarPC`. See details in :cite:`Brubeck2022` and :cite:`Brubeck2024`.
 :class:`.MassInvPC`
    Preconditioner for applying an inverse mass matrix.
 :class:`~.PCDPC`

firedrake/preconditioners/facet_split.py

@@ -19,7 +19,7 @@ class FacetSplitPC(PCBase):
     This allows for statically-condensed preconditioners to be applied to
     linear systems involving the matrix applied to the full set of DOFs. Code
     generated for the matrix-free operator evaluation in the space with full
-    DOFs will run faster than the one with interior-facet decoposition, since
+    DOFs will run faster than the one with interior-facet decomposition, since
     the full element has a simpler structure.
     """
 

firedrake/preconditioners/fdm.py

@@ -36,30 +36,6 @@
 import numpy
 import ctypes
 
-Citations().add("Brubeck2022", """
-@article{Brubeck2022,
-  title={A scalable and robust vertex-star relaxation for high-order {FEM}},
-  author={Brubeck, Pablo D. and Farrell, Patrick E.},
-  journal = {SIAM J. Sci. Comput.},
-  volume = {44},
-  number = {5},
-  pages = {A2991-A3017},
-  year = {2022},
-  doi = {10.1137/21M1444187}
-""")
-
-Citations().add("Brubeck2024", """
-@article{Brubeck2024,
-  title={{Multigrid solvers for the de Rham complex with optimal complexity in polynomial degree}},
-  author={Brubeck, Pablo D. and Farrell, Patrick E.},
-  journal = {SIAM J. Sci. Comput.},
-  volume = {46},
-  number = {3},
-  pages = {A1549-A1573},
-  year = {2024},
-  doi = {10.1137/22M1537370}
-""")
-
 
 __all__ = ("FDMPC", "PoissonFDMPC")
 

firedrake_citations/__init__.py

@@ -375,3 +375,29 @@ def print_at_exit(cls):
   title = {ngsPETSc: A coupling between NETGEN/NGSolve and PETSc},
   journal = {Journal of Open Source Software} }
 """)
+
+Citations().add("Brubeck2022", """
+@article{Brubeck2022,
+  author        = {Brubeck, Pablo D. and Farrell, Patrick E.},
+  doi           = {10.1137/21M1444187},
+  journal       = {SIAM J. Sci. Comput.},
+  number        = {5},
+  pages         = {A2991-A3017},
+  title         = {A scalable and robust vertex-star relaxation for high-order {FEM}},
+  volume        = {44},
+  year          = {2022}
+}
+""")
+
+Citations().add("Brubeck2024", """
+@article{Brubeck2024,
+  author        = {Brubeck, Pablo D. and Farrell, Patrick E.},
+  doi           = {10.1137/22M1537370},
+  journal       = {SIAM J. Sci. Comput.},
+  number        = {3},
+  pages         = {A1549-A1573},
+  title         = {{Multigrid solvers for the de Rham complex with optimal complexity in polynomial degree}},
+  volume        = {46},
+  year          = {2024}
+}
+""")

tests/firedrake/demos/test_demos_run.py

@@ -46,6 +46,7 @@
     Demo(("patch", "hdiv_riesz_star"), []),
     Demo(("quasigeostrophy_1layer", "qg_1layer_wave"), ["hypre", "vtk"]),
     Demo(("saddle_point_pc", "saddle_point_systems"), ["hypre", "mumps"]),
+    Demo(("fast_diagonalisation", "fast_diagonalisation_poisson"), ["mumps"]),
     Demo(('vlasov_poisson_1d', 'vp1d'), []),
 ]
 PARALLEL_DEMOS = [