Solving the Heat Equation with Irksome

Let’s start with the simple heat equation on \(\Omega = [0,10] \times [0,10]\), with boundary \(\Gamma\):

\[ \begin{align}\begin{aligned}u_t - \Delta u &= f\\u & = 0 \quad \textrm{on}\ \Gamma\end{aligned}\end{align} \]

for some known function \(f\). At each time \(t\), the solution to this equation will be some function \(u\in V\), for a suitable function space \(V\).

We transform this into weak form by multiplying by an arbitrary test function \(v\in V\) and integrating over \(\Omega\). We know have the variational problem of finding \(u:[0,T]\rightarrow V\) such that

\[(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

As usual, we need to import firedrake:

from firedrake import *

We will also need to import certain items from irksome:

from irksome import GaussLegendre, Dt, MeshConstant, TimeStepper

And we will need a little bit of UFL to support using the method of manufactured solutions:

from ufl.algorithms.ad import expand_derivatives

We will create the Butcher tableau for the lowest-order Gauss-Legendre Runge-Kutta method, which is more commonly known as the implicit midpoint rule:

butcher_tableau = GaussLegendre(1)
ns = butcher_tableau.num_stages

Now we define the mesh and piecewise linear approximating space in standard Firedrake fashion:

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

msh = RectangleMesh(N, N, x1, y1)
V = FunctionSpace(msh, "CG", 1)

We define variables to store the time step and current time value:

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

This defines the right-hand side using the method of manufactured solutions:

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))

We define the initial condition for the fully discrete problem, which will get overwritten at each time step:

u = Function(V)
u.interpolate(uexact)

Now, we will define the semidiscrete variational problem using standard UFL notation, augmented by the Dt operator from Irksome:

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")

Later demos will show how to use Firedrake’s sophisticated interface to PETSc for efficient block solvers, but for now, we will solve the system with a direct method (Note: the matrix type needs to be explicitly set to aij if sparse direct factorization were used with a multi-stage method, as we wind up with a mixed problem that Firedrake will assemble into a PETSc MatNest otherwise):

luparams = {"mat_type": "aij",
            "ksp_type": "preonly",
            "pc_type": "lu"}

Most of Irksome’s magic happens in the TimeStepper. It transforms our semidiscrete form F into a fully discrete form for the stage unknowns and sets up a variational problem to solve for the stages at each time step.:

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

This logic is pretty self-explanatory. We use the TimeStepper’s advance method, which solves the variational problem to compute the Runge-Kutta stage values and then updates the solution.:

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

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

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