Overview of the Devito domain specific language

from sympy import *
from devito import *

From equations to code in a few lines of Python – the main objective of this notebook is to demonstrate how Devito and its SymPy-powered symbolic API can be used to solve partial differential equations using the finite difference method with highly optimized stencils in a few lines of Python.

Defining the physical domain

A Grid object stores, among other things:

  • the physical extent (the size) of our domain, and
  • how many points we want to use in each dimension to discretise our data.

grid = Grid(shape=(5, 6), extent=(1., 1.))
grid

Functions, data, and expressions

To express our equation in symbolic form and discretise it using finite differences, Devito provides a set of Function types. A Function object also carries data.

f = Function(name='f', grid=grid)
f
f.data

By default, Devito Function objects use the spatial dimensions (x, y) for 2D grids and (x, y, z) for 3D grids. To solve a PDE over several timesteps a time dimension is also required by our symbolic function. For this Devito provides an additional function type, the TimeFunction, which incorporates the correct dimension along with some other intricacies needed to create a time stepping scheme.

g = TimeFunction(name='g', grid=grid)
g

Since the default time order of a TimeFunction is 1, the shape of f is (2, 5, 6), i.e. Devito has allocated two buffers to represent g(t, x, y) and g(t + dt, x, y):

g.shape

We can also create Function objects with custom Dimension’s.

x, y = grid.dimensions
d = Dimension(name='d')
u1 = Function(name='u', dimensions=(d, x, y), shape=(3,) + grid.shape)
u1
u2 = Function(name='u', dimensions=(y, x, d), shape=(6, 5, 3))
u2

Function’s are used to construct expressions. There is virtually no limit to the complexity an expression can have, but there’s a rule – it must be possible to construct an ordering of Dimension’s. In practice, this is never an issue.

cos(g)*f + sin(u1)  # OK, Devito can compile this expression
cos(g)*f + sin(u2)  # Not OK, Devito will complain because it sees both `x, y` and `y, x` as Function dimensions

Derivatives of symbolic functions

Devito provides a set of shorthand expressions (implemented as Python properties) that allow us to generate finite differences in symbolic form. For example, the property f.dx denotes \(\frac{\partial}{\partial x} f(x, y)\) - only that Devito has already discretised it with a finite difference expression.

f.dx

We can express derivatives of arbitrary order, but for this we’ll need to define a Function with a suitable spatial order. For example, the shorthand for the second derivative in x is .dx2, for the third order derivative .dx3, and so on.

h = Function(name='h', grid=grid, space_order=2)
h.dx2

We may also want to take a look at the stencil Devito will generate based on the chosen discretisation.

f.dx.evaluate
h.dx2.evaluate

A similar set of expressions exist for each spatial dimension defined on our grid, for example f.dy and f.dyl (here the l represents the left derivative). Obviously, one can also take derivatives in time of TimeFunction objects. For example, to take the first derivative in time of g you can simply write:

g.dt

There also exist convenient shortcuts to express the forward and backward stencil points, g(t+dt, x, y) and g(t-dt, x, y).

g.forward
g.backward

And of course, there’s nothing to stop us taking derivatives on these objects:

g.forward.dt
g.forward.dy

There also are shortcuts for classic differential operators

h.laplace
h.dx2 + h.dy2  # Equivalent to h.laplace

Some advanced features

More generally, we can take derivatives of arbitrary expressions

(g.dt + h.laplace + f.dy).dx2

Which can, depending on the chosen discretisation, lead to fairly complex stencils:

(g.dt + h.laplace + f.dy).dx2.evaluate

The DSL also extends naturally to tensorial objects

A = TensorFunction(name='A', grid=grid, space_order=2)
A
v = VectorFunction(name='v', grid=grid, space_order=2) 
v
b = A*v
b
div(b)

A linear convection operator

Note: The following example is derived from step 5 in the excellent tutorial series CFD Python: 12 steps to Navier-Stokes.

In this simple example we will show how to derive a very simple convection operator from a high-level description of the governing equation.

The governing equation we want to implement is the linear convection equation. We start off defining some parameters, such as the computational grid. We also initialize our velocity u with a smooth field:

from examples.cfd import init_smooth, plot_field

grid = Grid(shape=(81, 81), extent=(2., 2.))
u = TimeFunction(name='u', grid=grid, space_order=8)

# We can now set the initial condition and plot it
init_smooth(field=u.data[0], dx=grid.spacing[0], dy=grid.spacing[1])
init_smooth(field=u.data[1], dx=grid.spacing[0], dy=grid.spacing[1])

plot_field(u.data[0])

In particular, the linear convection equation that we want to implement is

\[\frac{\partial u}{\partial t}+c\frac{\partial u}{\partial x} + c\frac{\partial u}{\partial y} = 0.\]

Using the Devito shorthand notation, we can express the governing equations as:

c = 1.  # Value for c

eq = Eq(u.dt + c * u.dxl + c * u.dyl, 0)

We now need to rearrange our equation so that the term u(t+dt, x, y) is on the left-hand side, since it represents the next point in time for our state variable \(u\). Here, we use the Devito built-in solve function to create a valid stencil for our update to u(t+dt, x, y):

update = Eq(u.forward, solve(eq, u.forward))
update

Once we have created this update expression, we can create a Devito Operator. This Operator will basically behave like a Python function that we can call to apply the created stencil over our associated data.

op = Operator(update)  # Triggers compilation into C !
nt = 100  # Number of timesteps
dt = 0.2 * 2. / 80  # Timestep size (sigma=0.2)

op(t=nt+1, dt=dt)
plot_field(u.data[0])

Note that the real power of Devito is hidden within Operator, it will automatically generate and compile the optimized C code. We can look at this code (noting that this is not a requirement of executing it) via:

print(op.ccode)

What is not covered by this notebook

  • Mechanisms to expression injection and interpolation at grid points (“sparse operations”)
  • Subdomains and Conditionals
  • Boundary conditions (w/ and w/o subdomains)
  • Custom stencil coefficients
  • Staggered grids
  • …

How do I get parallel code?

op = Operator(update, language='openmp')
print(op)
op = Operator(update, language='openacc', platform='nvidiaX')
print(op)
Back to top