Following the initial convection tutorial with a single state variable $u$, we will now look at non-linear convection (step 6 in the original). This brings one new crucial challenge: computing a pair of coupled equations and thus updating two time-dependent variables $u$ and $v$.
The full set of coupled equations is now
\begin{aligned} \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} + v \frac{\partial u}{\partial y} = 0 \\ \\ \frac{\partial v}{\partial t} + u \frac{\partial v}{\partial x} + v \frac{\partial v}{\partial y} = 0\\ \end{aligned}and rearranging the discretized version gives us an expression for the update of both variables
\begin{aligned} u_{i,j}^{n+1} &= u_{i,j}^n - u_{i,j} \frac{\Delta t}{\Delta x} (u_{i,j}^n-u_{i-1,j}^n) - v_{i,j}^n \frac{\Delta t}{\Delta y} (u_{i,j}^n-u_{i,j-1}^n) \\ \\ v_{i,j}^{n+1} &= v_{i,j}^n - u_{i,j} \frac{\Delta t}{\Delta x} (v_{i,j}^n-v_{i-1,j}^n) - v_{i,j}^n \frac{\Delta t}{\Delta y} (v_{i,j}^n-v_{i,j-1}^n) \end{aligned}So, for starters we will re-create the original example run in pure NumPy array notation, before demonstrating the Devito version. Let's start again with some utilities and parameters:
In [1]:
from examples.cfd import plot_field, init_hat
import numpy as np
import sympy
%matplotlib inline
# Some variable declarations
nx = 101
ny = 101
nt = 80
c = 1.
dx = 2. / (nx - 1)
dy = 2. / (ny - 1)
sigma = .2
dt = sigma * dx
Let's re-create the initial setup with a 2D "hat function", but this time for two state variables.
In [2]:
#NBVAL_IGNORE_OUTPUT
# Allocate fields and assign initial conditions
u = np.empty((nx, ny))
v = np.empty((nx, ny))
init_hat(field=u, dx=dx, dy=dy, value=2.)
init_hat(field=v, dx=dx, dy=dy, value=2.)
plot_field(u)
Now we can create the two stencil expression for our two coupled equations according to the discretized equation above. We again use some simple Dirichlet boundary conditions to keep the values on all sides constant.
In [3]:
#NBVAL_IGNORE_OUTPUT
for n in range(nt + 1): ##loop across number of time steps
un = u.copy()
vn = v.copy()
u[1:, 1:] = (un[1:, 1:] -
(un[1:, 1:] * c * dt / dx * (un[1:, 1:] - un[1:, :-1])) -
vn[1:, 1:] * c * dt / dy * (un[1:, 1:] - un[:-1, 1:]))
v[1:, 1:] = (vn[1:, 1:] -
(un[1:, 1:] * c * dt / dx * (vn[1:, 1:] - vn[1:, :-1])) -
vn[1:, 1:] * c * dt / dy * (vn[1:, 1:] - vn[:-1, 1:]))
u[0, :] = 1
u[-1, :] = 1
u[:, 0] = 1
u[:, -1] = 1
v[0, :] = 1
v[-1, :] = 1
v[:, 0] = 1
v[:, -1] = 1
plot_field(u)
Excellent, we again get a wave that resembles the one from the oiginal examples.
Now we can set up our coupled problem in Devito. Let's start by creating two initial state variables $u$ and $v$, as before, and intialising them with our "hat function.
In [4]:
#NBVAL_IGNORE_OUTPUT
from devito import Grid, TimeFunction
# First we need two time-dependent data fields, both initialized with the hat function
grid = Grid(shape=(nx, ny), extent=(2., 2.))
u = TimeFunction(name='u', grid=grid)
init_hat(field=u.data[0], dx=dx, dy=dy, value=2.)
v = TimeFunction(name='v', grid=grid)
init_hat(field=v.data[0], dx=dx, dy=dy, value=2.)
plot_field(u.data[0])
Using the two TimeFunction
objects we can again derive our dicretized equation, rearrange for the forward stencil point in time and define our variable update expression - only we have to do everything twice now! We again use foward differences for time via u.dt
and backward differences in space via u.dxl
and u.dyl
to match the original tutorial.
In [5]:
from devito import Eq, solve
eq_u = Eq(u.dt + u*u.dxl + v*u.dyl)
eq_v = Eq(v.dt + u*v.dxl + v*v.dyl)
# We can use the same SymPy trick to generate two
# stencil expressions, one for each field update.
stencil_u = solve(eq_u, u.forward)
stencil_v = solve(eq_v, v.forward)
update_u = Eq(u.forward, stencil_u, subdomain=grid.interior)
update_v = Eq(v.forward, stencil_v, subdomain=grid.interior)
print("U update:\n%s\n" % update_u)
print("V update:\n%s\n" % update_v)
We then set Dirichlet boundary conditions at all sides of the domain to $1$.
In [6]:
x, y = grid.dimensions
t = grid.stepping_dim
bc_u = [Eq(u[t+1, 0, y], 1.)] # left
bc_u += [Eq(u[t+1, nx-1, y], 1.)] # right
bc_u += [Eq(u[t+1, x, ny-1], 1.)] # top
bc_u += [Eq(u[t+1, x, 0], 1.)] # bottom
bc_v = [Eq(v[t+1, 0, y], 1.)] # left
bc_v += [Eq(v[t+1, nx-1, y], 1.)] # right
bc_v += [Eq(v[t+1, x, ny-1], 1.)] # top
bc_v += [Eq(v[t+1, x, 0], 1.)] # bottom
And finally we can put it all together to build an operator and solve our coupled problem.
In [7]:
#NBVAL_IGNORE_OUTPUT
from devito import Operator
# Reset our data field and ICs
init_hat(field=u.data[0], dx=dx, dy=dy, value=2.)
init_hat(field=v.data[0], dx=dx, dy=dy, value=2.)
op = Operator([update_u, update_v] + bc_u + bc_v)
op(time=nt, dt=dt)
plot_field(u.data[0])
Excellent, we have now a scalar implementation of a convection problem, but this can be writen as a single vectorial equation:
$\frac{d U}{dt} + \nabla(U)U = 0$
Let's now use devito vectorial utilities and implement the vectorial equation
In [8]:
from devito import VectorTimeFunction, grad
U = VectorTimeFunction(name='U', grid=grid)
init_hat(field=U[0].data[0], dx=dx, dy=dy, value=2.)
init_hat(field=U[1].data[0], dx=dx, dy=dy, value=2.)
plot_field(U[1].data[0])
eq_u = Eq(U.dt + grad(U)*U)
We now have a vectorial equation. Unlike hte previous case, we do not need to play woth legt/right derivatives as the automated staggering of vectorial function takes care of it.
In [9]:
eq_u
Out[9]:
Then we set the nboundary conditions
In [10]:
x, y = grid.dimensions
t = grid.stepping_dim
bc_u = [Eq(U[0][t+1, 0, y], 1.)] # left
bc_u += [Eq(U[0][t+1, nx-1, y], 1.)] # right
bc_u += [Eq(U[0][t+1, x, ny-1], 1.)] # top
bc_u += [Eq(U[0][t+1, x, 0], 1.)] # bottom
bc_v = [Eq(U[1][t+1, 0, y], 1.)] # left
bc_v += [Eq(U[1][t+1, nx-1, y], 1.)] # right
bc_v += [Eq(U[1][t+1, x, ny-1], 1.)] # top
bc_v += [Eq(U[1][t+1, x, 0], 1.)] # bottom
In [11]:
# We can use the same SymPy trick to generate two
# stencil expressions, one for each field update.
stencil_U = solve(eq_u, U.forward)
update_U = Eq(U.forward, stencil_U, subdomain=grid.interior)
And we have the update (stencil) as a vectorial equation once again
In [12]:
update_U
Out[12]:
We finally run the operator
In [13]:
#NBVAL_IGNORE_OUTPUT
op = Operator([update_U] + bc_u + bc_v)
op(time=nt, dt=dt)
# The result is indeed the expected one.
plot_field(U[0].data[0])
In [14]:
from devito import norm
assert np.isclose(norm(u), norm(U[0]), rtol=1e-5, atol=0)