Accessing DIRK methods

Many practitioners favor diagonally implicit methods over fully implicit ones since the stages can be computed sequentially rather than concurrently. We support a range of DIRK methods and provide a convenient high-level interface that is very similar to other RK schemes. This demo is intended to show how to access DIRK methods seamlessly in Irksome.

This example uses the Qin Zhang symplectic DIRK to attack the mixed form of the wave equation. Let \(\Omega\) be the unit square with boundary \(\Gamma\). We write the wave equation as a first-order system of PDE:

\[ \begin{align}\begin{aligned}u_t + \nabla p & = 0\\p_t + \nabla \cdot u & = 0\end{aligned}\end{align} \]

together with homogeneous Dirichlet boundary conditions

\[p = 0 {\quad} \textrm{on}\ \Gamma\]

In this form, at each time, \(u\) is a vector-valued function in the Sobolev space \(H(\mathrm{div})\) and p is a scalar-valued function. If we select appropriate test functions \(v\) and \(w\), then we can arrive at the weak form (see the mixed wave demos for more information):

\[ \begin{align}\begin{aligned}(u_t, v) - (p, \nabla \cdot v) & = 0\\(p_t, w) + (\nabla \cdot u, w) & = 0\end{aligned}\end{align} \]

As in that case, we will use the next-to-lowest order Raviart-Thomas space for \(u\) and discontinuous piecewise linear elements for \(p\).

As an example, we will use the two-stage A-stable and symplectic DIRK of Qin and Zhang, given by Butcher tableau:

\[\begin{split}\begin{array}{c|cc} 1/4 & 1/4 & 0 \\ 3/4 & 1/2 & 1/4 \\ \hline & 1/2 & 1/2 \end{array}\end{split}\]

Imports from Firedrake and Irksome:

from firedrake import *
from irksome import QinZhang, Dt, MeshConstant, TimeStepper

We configure the discretization:

N = 10
msh = UnitSquareMesh(N, N)

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

V = FunctionSpace(msh, "RT", 2)
W = FunctionSpace(msh, "DG", 1)
Z = V*W

v, w = TestFunctions(Z)

butcher_tableau = QinZhang()

And set up the initial condition and variational problem:

x, y = SpatialCoordinate(msh)
up0 = project(as_vector([0, 0, sin(pi*x)*sin(pi*y)]), Z)

u0, p0 = split(up0)

F = inner(Dt(u0), v)*dx + inner(div(u0), w) * dx + inner(Dt(p0), w)*dx - inner(p0, div(v)) * dx

We will keep track of the energy to determine whether we’re sufficiently accurate in solving the linear system:

E = 0.5 * (inner(u0, u0)*dx + inner(p0, p0)*dx)

When stepping with a DIRK, we only solve for one stage at a time. Although we could configure PETSc to try some iterative solver, here we will just use a direct method:

params = {"mat_type": "aij",
          "snes_type": "ksponly",
          "ksp_type": "preonly",
          "pc_type": "lu"}

Now, we just set the stage type to be “dirk” in Irksome and we’re ready to advance in time:

stepper = TimeStepper(F, butcher_tableau, t, dt, up0,
                      stage_type="dirk",
                      solver_parameters=params)

print("Time    Energy")
print("==============")
while (float(t) < 1.0):
    if float(t) + float(dt) > 1.0:
        dt.assign(1.0 - float(t))

    stepper.advance()
    print("{0:1.1e} {1:5e}".format(float(t), assemble(E)))

    t.assign(float(t) + float(dt))

If all is right in the universe, you should see that the energy remains constant.