presentations/devito-dsl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview of the Devito domain specific language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import *\n",
    "from devito import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From equations to code in a few lines of Python -- the main objective of this notebook is to demonstrate how Devito and its [SymPy](http://www.sympy.org/en/index.html)-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.\n",
    "\n",
    "\n",
    "## Defining the physical domain\n",
    "\n",
    "A `Grid` object stores, among other things:\n",
    "\n",
    "* the physical `extent` (the size) of our domain, and\n",
    "* \n",
    "how many points we want to use in each dimension to discretise our data.\n",
    "\n",
    "<img src=\"figures/grid.png\" style=\"width: 220px;\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = Grid(shape=(5, 6), extent=(1., 1.))\n",
    "grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions, data, and expressions\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = Function(name='f', grid=grid)\n",
    "f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = TimeFunction(name='g', grid=grid)\n",
    "g"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also create `Function` objects with custom `Dimension`'s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = grid.dimensions\n",
    "d = Dimension(name='d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u1 = Function(name='u', dimensions=(d, x, y), shape=(3,) + grid.shape)\n",
    "u1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u2 = Function(name='u', dimensions=(y, x, d), shape=(6, 5, 3))\n",
    "u2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cos(g)*f + sin(u1)  # OK, Devito can compile this expression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cos(g)*f + sin(u2)  # Not OK, Devito will complain because it sees both `x, y` and `y, x` as Function dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Derivatives of symbolic functions\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f.dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = Function(name='h', grid=grid, space_order=2)\n",
    "h.dx2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We may also want to take a look at the stencil Devito will generate based on the chosen discretisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f.dx.evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h.dx2.evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.dt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There also exist convenient shortcuts to express the forward and backward stencil points, `g(t+dt, x, y)` and `g(t-dt, x, y)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.forward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.backward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And of course, there's nothing to stop us taking derivatives on these objects:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.forward.dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.forward.dy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There also are shortcuts for classic differential operators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h.laplace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h.dx2 + h.dy2  # Equivalent to h.laplace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Some advanced features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More generally, we can take **derivatives of arbitrary expressions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(g.dt + h.laplace + f.dy).dx2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which can, depending on the chosen discretisation, lead to fairly complex stencils: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(g.dt + h.laplace + f.dy).dx2.evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The DSL also extends naturally to **tensorial objects**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = TensorFunction(name='A', grid=grid, space_order=2)\n",
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = VectorFunction(name='v', grid=grid, space_order=2) \n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = A*v\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div(b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A linear convection operator\n",
    "\n",
    "**Note:** The following example is derived from [step 5](http://nbviewer.ipython.org/github/barbagroup/CFDPython/blob/master/lessons/07_Step_5.ipynb) in the excellent tutorial series [CFD Python: 12 steps to Navier-Stokes](http://lorenabarba.com/blog/cfd-python-12-steps-to-navier-stokes/).\n",
    "\n",
    "In this simple example we will show how to derive a very simple convection operator from a high-level description of the governing equation.\n",
    "\n",
    "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:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [],
   "source": [
    "from examples.cfd import init_smooth, plot_field\n",
    "\n",
    "grid = Grid(shape=(81, 81), extent=(2., 2.))\n",
    "u = TimeFunction(name='u', grid=grid, space_order=8)\n",
    "\n",
    "# We can now set the initial condition and plot it\n",
    "init_smooth(field=u.data[0], dx=grid.spacing[0], dy=grid.spacing[1])\n",
    "init_smooth(field=u.data[1], dx=grid.spacing[0], dy=grid.spacing[1])\n",
    "\n",
    "plot_field(u.data[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In particular, the linear convection equation that we want to implement is\n",
    "\n",
    "$$\\frac{\\partial u}{\\partial t}+c\\frac{\\partial u}{\\partial x} + c\\frac{\\partial u}{\\partial y} = 0.$$\n",
    "\n",
    "Using the Devito shorthand notation, we can express the governing equations as:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = 1.  # Value for c\n",
    "\n",
    "eq = Eq(u.dt + c * u.dxl + c * u.dyl, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "update = Eq(u.forward, solve(eq, u.forward))\n",
    "update"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [],
   "source": [
    "op = Operator(update)  # Triggers compilation into C !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = 100  # Number of timesteps\n",
    "dt = 0.2 * 2. / 80  # Timestep size (sigma=0.2)\n",
    "\n",
    "op(t=nt+1, dt=dt)\n",
    "plot_field(u.data[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [],
   "source": [
    "print(op.ccode)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is not covered by this notebook\n",
    "\n",
    "* Mechanisms to expression injection and interpolation at grid points (\"sparse operations\")\n",
    "* Subdomains and Conditionals\n",
    "* Boundary conditions (w/ and w/o subdomains)\n",
    "* Custom stencil coefficients\n",
    "* Staggered grids\n",
    "* ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How do I get parallel code?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "op = Operator(update, language='openmp')\n",
    "print(op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "op = Operator(update, language='openacc', platform='nvidiaX')\n",
    "print(op)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}