Form compilation

As we discussed in Code generation DOLFINx generates code once the user has specified the variational form. This process can be somewhat time consuming, as generating, compiling and linking the code can take time. We start by setting up a simple unit square and a first order Lagrange space.

import numpy as np
import ufl
from mpi4py import MPI
import dolfinx

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 30, 30)
V = dolfinx.fem.FunctionSpace(mesh, ("Lagrange", 1))

Let us consider a simple heat equation,

(6)\[\begin{align} \frac{\partial u}{\partial t} - \nabla \cdot( k(t) \nabla u) &= f(x,y,t) \qquad \text{in } \Omega \\ \frac{\partial u}{\partial n} &= 0 \qquad \text{on } \partial \Omega\\ u(\cdot, 0) &= \frac{1}{2\pi \sigma} e^{-\frac{1}{2}\left(\frac{x-\mu}{\sigma}\right)^2} e^{-\frac{1}{2}\left(\frac{y-\mu}{\sigma}\right)^2} \\ k(t) &= \begin{cases} 0.1 \quad \text{if } t<0.5\\ 0.05 \quad \text{if } t>=0.5 \end{cases}\\ f(x,y,t) &= \begin{cases} 0.4\cdot y \quad \text{if } x<0.5\\ 0.5\cdot t\quad \text{if } x>=0.5 \end{cases} \end{align}\]

In this equation, there are several time dependent and spatially varying coefficients. With a naive implementation of this code in DOLFINx, one would define the variational formulation inside the temporal loop, and re-compile the variational form for each time step. In DOLFINx, a form is compiled whenever one calls dolfinx.fem.form, or initialize a dolfinx.fem.petsc.LinearProblem. We start by defining the variational form, with a backward Euler time stepping.

Time-dependent constants

To be able to use adaptive time stepping, we define dt as a dolfinx.fem.Constant, such that we can re-assign values to the constant without having to re-compile the code.

dt = dolfinx.fem.Constant(mesh, 0.01)

It is easy to re-assign a value to the constant by calling

dt.value = 0.005

Similarly, we can define the diffusive coefficient k such as

def k_func(t):
    return 0.1 if t < 0.5 else 0.05

k = dolfinx.fem.Constant(mesh, 0.1)
t = dolfinx.fem.Constant(mesh, 0.)
while t.value < 1:
    # Update t
    t.value += dt.value
    # Update k
    k.value = k_func(t.value)
# We reset t
t.value = 0

Conditional values

Next, we define the spatial and temporal source term using a ufl.conditional We start by defining the spatially varying parameters of the mesh

x, y = ufl.SpatialCoordinate(mesh)

# Next, we create each component of the conditional statement
condition = ufl.lt(x, 0.5)
true_statement = 0.4 * y
false_statement = 0.5 * t
f = ufl.conditional(condition, true_statement, false_statement)

Inside the compile code, this condition will be evaluated at every quadrature point.

Now we have defined all varying functions outside the time loop, and each of them can be updated by updating the constant values

We define the full variational formulation

u = ufl.TrialFunction(V)
u_n = dolfinx.fem.Function(V)
dudt = (u - u_n) / dt

v = ufl.TestFunction(V)
dx = ufl.Measure("dx", domain=mesh)
F = dudt * v * dx + k * ufl.inner(ufl.grad(u), ufl.grad(v)) * dx - f * v * dx
a, L = ufl.system(F)

We generate and compile the C code for these expressions using dolfinx.fem.form

a_compiled = dolfinx.fem.form(a)
L_compiled = dolfinx.fem.form(L)

We generate the initial condition by using lambda expressions

def u_init(x, sigma=0.1, mu=0.3):
    """
    The input function x is a (3, number_of_points) numpy array, which is then
    evaluated with the vectorized numpy functions for efficiency
    """
    return 1./(2 * np.pi * sigma)*np.exp(-0.5*((x[0]-mu)/sigma)**2)*np.exp(-0.5*((x[1]-mu)/sigma)**2)

u_n.interpolate(u_init)

Setting up the linear solver

As we have now defined a and L, we are ready to set up a solver for the arising linear system. We use PETSc for this, and a direct solver.

uh = dolfinx.fem.Function(V)
petsc_options = {"ksp_type": "preonly",
                 "pc_type": "lu", "pc_factor_solver_type": "mumps"}
problem = dolfinx.fem.petsc.LinearProblem(
    a_compiled, L_compiled, u=uh, bcs=[], petsc_options=petsc_options)

For each temporal step, we update the time variable and call the solve command that re-assemble the system

T = 1
while t.value < T:
    t.value += dt.value
    k.value = k_func(t.value)
    problem.solve()
    # Update previous solution
    u_n.x.array[:] = uh.x.array