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 RadauIIA, 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 3-stage RadauIIA Runge-Kutta method, which has an embedded scheme we can use for adaptivity:

butcher_tableau = RadauIIA(3)
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, current time value, and final time:

MC = MeshConstant(msh)
dt = MC.Constant(10.0 / N)
t = MC.Constant(0.0)
Tf = MC.Constant(1.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.

In this demo, we use an adaptive timestepper, selected by sending a dictionary of adaptive parameters, which tells Irksome to also compute an error estimate at each timestep, and use that estimate to adjust the timestep size. The adaptation depends on some given parameters, including a tolerance tol for for the error at each step, and KI and KP that set the so-called integration and performance gain for the estimate.:

adapt_params = {"tol":1e-3, "KI":1/15, "KP":0.13}
stepper = TimeStepper(F, butcher_tableau, t, dt, u, bcs=bc,
                      stage_type="deriv", solver_parameters=luparams,
                      adaptive_parameters = adapt_params)

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.

Here, in contrast to the non-adaptive case, we get an estimate of the error at each step (that we do not use here) and a new adaptive timestep size at each step. We use these to control integrating to a fixed final time, Tf. This exposes the dt_max data for TimeStepper, which puts a hard limit on the timestep size in the adaptive case.:

while (float(t) < float(Tf)):
    stepper.dt_max = float(Tf)-float(t)
    (adapt_error, adapt_dt) = stepper.advance()
    print(float(t))
    t.assign(float(t) + float(adapt_dt))

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

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