Solving the Advection-Diffusion Equation

The Advection-Diffusion Equation models the distribution of a pollutant quantity $c$ in some medium (i.e. a fluid). It comprises two parts:

The diffusion term $\nabla \cdot (D \nabla c)$ describes the dispersion of the concentration $c$ due to molecular motion, according to the diffusivity $D$ ;
The advection term $\nabla \cdot (\boldsymbol{v} c)$ represents the motion of the concentration $c$ due to being carried with the medium’s velocity field $\boldsymbol{v}$ ;
The transient term $\frac{\partial c}{\partial t}$ relates the time evolution of $c$ to the other two spatial derivatives.

These two terms together describe the evolution in time of the concentration $c$ :

\frac{\partial c}{\partial t} = \nabla \cdot (D \nabla c) - \nabla \cdot (\boldsymbol{v} c).

(1)

This lesson is about using pystencils to simulate a simple two-dimensional advection-diffusion scenario. We will discretize the ADE using a second-order finite difference scheme in space, and first-order explicit Euler integration in time; set up boundary conditions and an integration loop, and visualize the results. As an extra perk, we will build the entire solver to be accellerated by your machine’s graphics processor (GPU) ^[1].

Lesson Goals

During this lesson, you will use pystencils to construct a solver for a simple two-dimensional, time-dependent problem governed by the advection-diffusion equation. In the process, you will

derive an integration rule for the ADE using second-order finite differences in space, and an explicit Euler scheme in time;
implement a Neumann boundary condition as a pystencils operator;
implement a localized mass source using a guarded operator;
accelerate your simulation by running it on the GPU;
and use PyVista to interactively visualize your simulation results.

Let’s start by importing the packages we will need.

import sympy as sp
import pystencils as ps
import numpy as np
import pyvista as pv

interactive_plots = False

GPU Target Selection¶

Up to this point, we have only generated non-optimized kernel code for CPU execution. Pystencils, however, is built to harness the parallel processing power that is now built into most computers, from your laptop up to exascale computing clusters.

One major driver of that parallelism are general-purpose graphics processing units (GPUs). If you have a GPU, you can use it to accelerate your pystencils operators with only minimal changes to your code - without the need to get into the details of GPU kernel programming.

Exercise 1 (Install and Check CuPy)

Pystencils relies on the Python library cupy for accessing Nvidia and AMD GPUs. Use the following code to check if cupy is installed in your environment and if it can access your GPU:

import cupy as cp

cp.cuda.runtime.getDeviceProperties(0)

If everything is set up correctly, this will show you the details of your system’s GPU.

If you do not have CuPy installed yet, follow the installation instructions in order to install the appropriate version for your system. When done, restart your Jupyter kernel and re-check the installation using the above code.

Create Algebraic Objects¶

Let’s begin building our ADE simulation by setting up algebraic objects for our simulation domain, variables, and fields.

L = (200, 80)
Ω = ps.grids.Patch("Ω", L, num_vertices=(lx + 1 for lx in L))
                  
D, Δx, Δt = sp.symbols("D, Δx, Δt")

c = ps.grids.TensorField("c", Ω.vertices, ())
v = ps.grids.TensorField("v", Ω.vertices, (2,))

Implement the Integrator¶

Next up, we implement the time-integration operator for the advection-diffusion equation. To do so, we first need to take a look at our discretization scheme. We can discretize the three derivative terms separately, like follows:

For the diffusion term $\nabla \cdot (D \nabla c)$ , we make the simplifying assumption that $D$ is constant across the domain. It thus simplifies to $D \Delta c$ with the Laplace operator $\Delta$ . To discretize, we apply the same second-order finite difference stencil as in the previous lesson to receive:
$D \left( \operatorname{Diff}_{\Delta x} \boldsymbol{c} \right)_{ij} = \frac{D}{{\Delta x}^2} \left( - 4 c_{i,j} + c_{i-1, j} + c_{i+1,j} + c_{i,j-1} + c_{i,j+1} \right)$
(2)
The advection term $\nabla \cdot (\boldsymbol{v} c)$ , signifies the divergence of the vector field $\boldsymbol{v} c$ :
$\nabla \cdot (\boldsymbol{v} c) = \frac{\partial}{\partial x} (v^x c) + \frac{\partial}{\partial y} (v^y c).$
(3)
To discretize, we simply apply central differences to the two spatial derivatives:
$\begin{align*} \left( \operatorname{Adv}_{\Delta x} (\boldsymbol{v} c) \right)_{ij} = & \frac{1}{2 \Delta x} \left( c_{i+1, j} v^x_{i+1,j} - c_{i-1, j} v^x_{i-1,j} \right) \\ & + \frac{1}{2 \Delta x} \left( c_{i, j+1} v^y_{i,j+1} - c_{i, j-1} v^y_{i,j-1} \right) \end{align*}$
(4)
The transient term $\frac{\partial c}{\partial t}$ we discretize via a forward-Euler difference:
$\left( \frac{\partial c}{\partial t} \right)_{ij} = \frac{ c_{ij}^{(t + 1)} - c_{ij}^{(t)} }{ \Delta t }$
(5)
By rewriting (1), this gives rise to the following explicit update scheme:
$c^{(t + 1)} = c^{(t)} + {\Delta t} \left( D \operatorname{Diff}_{\Delta x} (\boldsymbol{c}) - \operatorname{Adv}_{\Delta x} (\boldsymbol{v} c) \right)$
(6)

c_next = ps.grids.TensorField("c_next", Ω.vertices, ())


@ps.flow.operator(
    target=target,
    ispace=Ω.vertices[1:-1, 1:-1]
)
def integrate_ade(_eq):
    _eq.let[Δx] = Ω.spacing[0]

    diff, advc = sp.symbols("diff, advc")
    _eq.let[diff] = D * (
                          c[-1, 0]() 
                        + c[1, 0]()
                        + c[0, -1]()
                        + c[0, 1]()
                        - 4*c[0, 0]()
                    ) / Δx**2
    _eq.let[advc] = (
                        c[1, 0]()*v[1, 0](0) - c[-1, 0]()*v[-1, 0](0) 
                        + c[0, 1]()*v[0, 1](1) - c[0, -1]()*v[0, -1](1)
                    ) / (2*Δx)

    transient = (c_next() - c()) / Δt
    euler_update = sp.solve(transient - diff + advc, c_next())[0]
    _eq.store[c_next()] = euler_update

Prepare the Simulation Scenario¶

Our next step is to set up the simulation data. This includes allocating the field arrays, as well as initializing the field $\boldsymbol{v}$ to a velocity profile along which the quantity $c$ will be advected.

If we want to run our simulation on the GPU, the data arrays must accordingly be allocated in GPU memory. The PatchData class takes care of this automatically if we tell it which target we are running on.

Initialize the Velocity Field¶

@ps.flow.operator(target=target)
def set_v(_eq):
    y, t, r, v_max = sp.symbols("y, t, r, v_max")
    _eq.let[y] = Ω.vertex()[1]
    _eq.let[r] = Ω.extents[1] / 2
    _eq.let[t] = 1 - ((y - r) / r)**2
    _eq.let[v_max] = 2

    _eq.store[v(0)] = t * v_max
    _eq.store[v(1)] = 0

Run the following cell to initialize the velocity field:

set_v(Ω_data)

Use PyVista for Visualization¶

In the following cell, we use PyVista to visualize the $x$ -component of $\boldsymbol{v}$ .

When PyVista is installed, you can use the viz() method of a patch data object to get its visualization bridge. You can then obtain a VTK mesh object that represents the patch’s simulation grid, and contains the data of its fields. Pass that mesh to a PyVista Plotter and use it to show the visualization:

<PIL.Image.Image image mode=RGB size=1024x768> — $x$ -component of the velocity field

Boundary Conditions and a Mass Source¶

Apart from the time integrator for the ADE, our simulation will still need boundary conditions at its walls. Also, for anything at all to happen, we would like to inject a fixed amount of the concentration $c$ into the medium at a certain point.

Define the Boundary Conditions¶

We would like to use two different kinds of boundary conditions in our simulation:

On the west, north, and south borders of $\Omega$ , we apply a homogenous Dirichlet boundary condition; that is, we want to set $c_{ij} = 0$ at all times.
At the eastern border, we are going to place an outflow where the concentration $c$ can leave the domain. We model this using a Neumann boundary condition:
$\frac{\partial c}{\partial x} = 0.$
(8)

Interestingly, we can realize the Dirichlet boundaries simply by doing nothing at all. The PatchData has initialized the arrays for c and c_next to zero; and if you have correctly configured the iteration region of your time integrator, the values at the boundary will stay at zero for all time.

Inject a Mass Source¶

Last, to complete our simulation setup, we would like to inject some concentration $c$ into our system using a circular mass source. We do so by fixing the concentration to $c(x, y) = 1.0$ within the circle with radius $r = L_y / 10$ around the point $\boldsymbol{p} = \begin{pmatrix} L_x / 4 \\ L_y / 2 \end{pmatrix}$ . This leads us to a situation where we need to write an operator that performs this operation only on an irregularily shaped (circular) part of the domain. Restricting the operator’s iteration space to that circle using ispace is rather difficult, so we instead use another way: we are going to construct a guarded operator.

Pystencils allows us to guard parts of an operator according to given conditions; they will only be executed when these conditions evaluate to true. In our case, our operator has two parts:

The equations to compute the distance $d$ of the current grid node to the circle’s center point $\boldsymbol{p}$ ,
and the guarded portion that writes $c = 1.0$ only if $d < r$ .

We realize these parts by splitting our operator into blocks, and by connecting these blocks accordingly. We use the ps.flow.block and ps.flow.guarded_block decorators to create regular and guarded blocks, respectively.

Exercise 8 (ADE Problem: Mass Source)

Complete the following implementation of the constant_source operator by filling in the gaps (marked by ...):

Complete the block compute_distance such that the symbol d is set to the distance of the current vertex to the circle’s center point
Fill in the condition under which the guarded block inject_source should be executed
Complete inject_source to set $c_{ij}$ to 1

@ps.flow.operator(target=target)
def constant_source():
    d = sp.Symbol("d")

    @ps.flow.block
    def compute_distance(_eq):
        ...
        _eq.export[d] = ...

    @ps.flow.guarded_block( ...condition... , preds=[compute_distance])
    def inject_source(_eq):
        ...

    return inject_source

On Connecting Blocks

When building operators out of multiple blocks, take care of the following:

Blocks need to be connected to form a directed acyclic graph. This is done by adding earlier blocks as predecessors to later blocks using the preds= keyword argument.
When using multiple blocks, use export instead of let to define equations for symbols that should be visible in successor blocks. In the above operator, _eq.export[d] = makes sure that the definition of d is visible in inject_source.
End your operator with return <end-block> to give the operator’s last “output” block back to pystencils; the operator system will then collect all blocks connected to it.

@ps.flow.operator(target=target)
def constant_source():
    x, y, cx, cy, radius, d = sp.symbols("x, y, cx, cy, radius, d")

    @ps.flow.block
    def compute_distance(_eq):
        _eq.let[x, y] = Ω.vertex()
        _eq.let[cx, cy] = L[0] / 4, L[1] / 2
        _eq.export[radius] = L[1] / 10
        _eq.export[d] = sp.sqrt((x-cx)**2 + (y-cy)**2)

    @ps.flow.guarded_block(d < radius, preds=[compute_distance])
    def inject_source(_eq):
        _eq.store[c()] = 1

    return inject_source

Run the Simulation¶

Finally, all the ingredients are in place, and we can run our simulation. Let’s first define the full integration routine:

def integrate():
    neumann_outflow(Ω_data)
    constant_source(Ω_data)
    integrate_ade(Ω_data)
    Ω_data.swap(c, c_next)

The integration loop is built to run a given number nsteps of time steps, and to extract snapshots of the field $c$ at a given frequency, in order to later visualize the time series:

def simloop(nsteps, capture_freq=100):
    c_frames = []
    for t in range(nsteps):
        if t % capture_freq == 0:
            c_frames.append(Ω_data.asnumpy(c))
        integrate()

    return c_frames

Let’s run the simulation for 4000 time steps, and then visualize the result.

Ω_data[c][:] = 0.

c_frames = simloop(4000)

We again use PyVista to look at our simulation results. In the Jupyter notebook, you can use the following cell to create an interactive widget for exploring the captured time series:

If everything went well, you can see how a trail of advected material develops in the wake of the mass source. The larger the diffusivity $D$ , the wider the trail will grow. You can now experiment by running the simulation with different values for the diffusivity, but also different time intervals $\Delta t$ .

Conclusion¶

This concludes our lesson on implementing an explicit solver for the ADE in pystencils. We have seen how to configure operators to run on the GPU; how to bridge the patch data to PyVista for visualizaton, and how to use guarded kernels to apply operations only in certain regions of our domain. The end result is your first simulation of a time-dependent problem in pystencils; we will see many more of these in the following lessons.

Footnotes¶

If you have one.
↩