Solving the heat equation with monolithic multigrid

This reprise of the heat equation demo uses a monolithic multigrid algorithm suggested by Patrick Farrell to perform time advancement.

We consider the heat equation on \(\Omega = [0,10] \times [0,10]\), with boundary \(\Gamma\): giving rise to the weak form

\[(u_t, v) + (\nabla u, \nabla v) = (f, v)\]

This demo implements an example used by Solin with a particular choice of \(f\) given below

We perform similar imports and setup as before:

from firedrake import *
from irksome import GaussLegendre, Dt, MeshConstant, TimeStepper
from ufl.algorithms.ad import expand_derivatives
butcher_tableau = GaussLegendre(2)

However, we need to set up a mesh hierarchy to enable geometric multigrid within Firedrake:

N = 128
x0 = 0.0
x1 = 10.0
y0 = 0.0
y1 = 10.0

from math import log
coarseN = 8  # size of coarse grid
nrefs = log(N/coarseN, 2)
assert nrefs == int(nrefs)
nrefs = int(nrefs)
base = RectangleMesh(coarseN, coarseN, x1, y1)
mh = MeshHierarchy(base, nrefs)
msh = mh[-1]

From here, setting up the function space, manufactured solution, etc, are just as for the regular heat equation demo:

V = FunctionSpace(msh, "CG", 1)

MC = MeshConstant(msh)
dt = MC.Constant(10.0 / N)
t = MC.Constant(0.0)

x, y = SpatialCoordinate(msh)
S = Constant(2.0)
C = Constant(1000.0)
B = (x-Constant(x0))*(x-Constant(x1))*(y-Constant(y0))*(y-Constant(y1))/C
R = (x * x + y * y) ** 0.5
uexact = B * atan(t)*(pi / 2.0 - atan(S * (R - t)))
rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact))

u = Function(V)
u.interpolate(uexact)
v = TestFunction(V)
F = inner(Dt(u), v)*dx + inner(grad(u), grad(v))*dx - inner(rhs, v)*dx
bc = DirichletBC(V, 0, "on_boundary")

And now for the solver parameters. Note that we are solving a block-wise system with all stages coupled together. This performs a monolithic multigrid with pointwise block Jacobi preconditioning:

mgparams = {"mat_type": "aij",
            "snes_type": "ksponly",
            "ksp_type": "gmres",
            "ksp_monitor_true_residual": None,
            "pc_type": "mg",
            "mg_levels": {
                "ksp_type": "chebyshev",
                "ksp_max_it": 1,
                "ksp_convergence_test": "skip",
                "pc_type": "python",
                "pc_python_type": "firedrake.PatchPC",
                "patch": {
                    "pc_patch": {
                        "save_operators": True,
                        "partition_of_unity": False,
                        "construct_type": "star",
                        "construct_dim": 0,
                        "sub_mat_type": "seqdense",
                        "dense_inverse": True,
                        "precompute_element_tensors": None},
                    "sub": {
                        "ksp_type": "preonly",
                        "pc_type": "lu"}}},
            "mg_coarse": {
                "pc_type": "lu",
                "pc_factor_mat_solver_type": "mumps"}
            }

These solver parameters work just fine in the TimeStepper:

stepper = TimeStepper(F, butcher_tableau, t, dt, u, bcs=bc,
                      solver_parameters=mgparams)

And we can advance the solution in time in typical fashion:

while (float(t) < 1.0):
    if (float(t) + float(dt) > 1.0):
        dt.assign(1.0 - float(t))
    stepper.advance()
    print(float(t), flush=True)
    t.assign(float(t) + float(dt))

After the solve, we can retrieve some statistics about the solver:

steps, nonlinear_its, linear_its = stepper.solver_stats()

print("Total number of timesteps was %d" % (steps))
print("Average number of nonlinear iterations per timestep was %.2f" % (nonlinear_its/steps))
print("Average number of linear iterations per timestep was %.2f" % (linear_its/steps))

Finally, we print out the relative \(L^2\) error:

print()
print(norm(u-uexact)/norm(uexact))