# Chapter 11: Partial differential equations¶

Robert Johansson

Source code listings for Numerical Python - A Practical Techniques Approach for Industry (ISBN 978-1-484205-54-9).

The source code listings can be downloaded from http://www.apress.com/9781484205549

In [1]:
import numpy as np

In [2]:
%matplotlib inline
%config InlineBackend.figure_format='retina'
import matplotlib.pyplot as plt
import matplotlib as mpl

In [3]:
import mpl_toolkits.mplot3d

In [4]:
import scipy.sparse as sp

In [5]:
import scipy.sparse.linalg

In [6]:
import scipy.linalg as la


## 1d example¶

Heat equation:

$$-5 = u_{xx}, u(x=0) = 1, u(x=1) = 2$$

$$u_{xx}[n] = (u[n-1] - 2u[n] + u[n+1])/dx^2$$

In [7]:
N = 5

In [8]:
u0 = 1
u1 = 2

In [9]:
dx = 1.0 / (N + 1)

In [10]:
A = (np.eye(N, k=-1) - 2 * np.eye(N) + np.eye(N, k=1)) / dx**2

In [11]:
A

Out[11]:
array([[-72.,  36.,   0.,   0.,   0.],
[ 36., -72.,  36.,   0.,   0.],
[  0.,  36., -72.,  36.,   0.],
[  0.,   0.,  36., -72.,  36.],
[  0.,   0.,   0.,  36., -72.]])
In [12]:
d = -5 * np.ones(N)
d[0] -= u0 / dx**2
d[N-1] -= u1 / dx**2

In [13]:
u = np.linalg.solve(A, d)

In [14]:
x = np.linspace(0, 1, N+2)

In [15]:
U = np.hstack([[u0], u, [u1]])

In [16]:
fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(x, U)
ax.plot(x[1:-1], u, 'ks')
ax.set_xlim(0, 1)
ax.set_xlabel(r"$x$", fontsize=18)
ax.set_ylabel(r"$u(x)$", fontsize=18)
fig.savefig("ch11-fdm-1d.pdf")
fig.tight_layout();


## 2d example¶

laplace equation: $u_{xx} + u_{yy} = 0$

on boundary:

$$u(x=0) = u(x=1) = u(y = 0) = u(y = 1) = 10$$

$$u_{xx}[m, n] = (u[m-1, n] - 2u[m,n] + u[m+1,n])/dx^2$$

$$u_{yy}[m, n] = (u[m, n-1] - 2u[m,n] + u[m,n+1])/dy^2$$

final equation

$$# 0 ¶ (u[m-1 + N n] - 2u[m + N n] + u[m+1 + N n])/dx^2 + # (u[m + N (n-1)] - 2u[m + N n] + u[m + N(n+1]))/dy^2¶ (u[m + N n -1] - 2u[m + N n] + u[m + N n + 1])/dx^2 + (u[m + N n -N)] - 2u[m + N n] + u[m + N n + N]))/dy^2$$

In [17]:
N = 100

In [18]:
u0_t, u0_b = 5, -5

In [19]:
u0_l, u0_r = 3, -1

In [20]:
dx = 1. / (N+1)

In [21]:
A_1d = (sp.eye(N, k=-1) + sp.eye(N, k=1) - 4 * sp.eye(N))/dx**2

In [22]:
A = sp.kron(sp.eye(N), A_1d) + (sp.eye(N**2, k=-N) + sp.eye(N**2, k=N))/dx**2

In [23]:
A

Out[23]:
<10000x10000 sparse matrix of type '<type 'numpy.float64'>'
with 49600 stored elements in Compressed Sparse Row format>
In [24]:
A.nnz * 1.0/ np.prod(A.shape) * 2000

Out[24]:
0.992
In [25]:
d = np.zeros((N, N))

d[0, :] += -u0_b
d[-1, :] += -u0_t
d[:, 0] += -u0_l
d[:, -1] += -u0_r

d = d.reshape(N**2) / dx**2

In [26]:
u = sp.linalg.spsolve(A, d).reshape(N, N)

In [27]:
U = np.vstack([np.ones((1, N+2)) * u0_b,
np.hstack([np.ones((N, 1)) * u0_l, u, np.ones((N, 1)) * u0_r]),
np.ones((1, N+2)) * u0_t])

In [28]:
fig, ax = plt.subplots(1, 1, figsize=(8, 6))

x = np.linspace(0, 1, N+2)
X, Y = np.meshgrid(x, x)

c = ax.pcolor(X, Y, U, vmin=-5, vmax=5, cmap=mpl.cm.get_cmap('RdBu_r'))
cb = plt.colorbar(c, ax=ax)

ax.set_xlabel(r"$x_1$", fontsize=18)
ax.set_ylabel(r"$x_2$", fontsize=18)
cb.set_label(r"$u(x_1, x_2)$", fontsize=18)
fig.savefig("ch11-fdm-2d.pdf")
fig.tight_layout()

In [29]:
x = np.linspace(0, 1, N+2)
X, Y = np.meshgrid(x, x)

In [30]:
fig = plt.figure(figsize=(12, 5.5))
cmap = mpl.cm.get_cmap('RdBu_r')

ax = fig.add_subplot(1, 2, 1)
p = ax.pcolor(X, Y, U, vmin=-5, vmax=5, cmap=cmap)
ax.set_xlabel(r"$x_1$", fontsize=18)
ax.set_ylabel(r"$x_2$", fontsize=18)

ax = fig.add_subplot(1, 2, 2, projection='3d')
p = ax.plot_surface(X, Y, U, vmin=-5, vmax=5, rstride=3, cstride=3, linewidth=0, cmap=cmap)
ax.set_xlabel(r"$x_1$", fontsize=16)
ax.set_ylabel(r"$x_2$", fontsize=16)
cb = plt.colorbar(p, ax=ax, shrink=0.75)
cb.set_label(r"$u(x_1, x_2)$", fontsize=18)

fig.savefig("ch11-fdm-2d.pdf")
fig.savefig("ch11-fdm-2d.png")
fig.tight_layout()


### Compare performance when using dense/sparse matrices¶

In [31]:
A_dense = A.todense()

In [32]:
%timeit np.linalg.solve(A_dense, d)

1 loops, best of 3: 10.1 s per loop

In [33]:
%timeit la.solve(A_dense, d)

1 loops, best of 3: 12.2 s per loop

In [34]:
%timeit sp.linalg.spsolve(A, d)

10 loops, best of 3: 33.9 ms per loop

In [35]:
10.8 / 31.9e-3

Out[35]:
338.5579937304076

### 2d example with source term¶

In [36]:
d = - np.ones((N, N))
d = d.reshape(N**2)

In [37]:
u = sp.linalg.spsolve(A, d).reshape(N, N)

In [38]:
U = np.vstack([np.zeros((1, N+2)),
np.hstack([np.zeros((N, 1)), u, np.zeros((N, 1))]),
np.zeros((1, N+2))])

In [39]:
x = np.linspace(0, 1, N+2)
X, Y = np.meshgrid(x, x)

In [40]:
fig, ax = plt.subplots(1, 1, figsize=(8, 6), subplot_kw={'projection': '3d'})

p = ax.plot_surface(X, Y, U, rstride=4, cstride=4, linewidth=0, cmap=mpl.cm.get_cmap("Reds"))
cb = fig.colorbar(p, shrink=0.5)

ax.set_xlabel(r"$x_1$", fontsize=18)
ax.set_ylabel(r"$x_2$", fontsize=18)
cb.set_label(r"$u(x_1, x_2)$", fontsize=18)


## FEniCS¶

In [41]:
import dolfin

In [42]:
import mshr

In [43]:
dolfin.parameters["reorder_dofs_serial"] = False
dolfin.parameters["allow_extrapolation"] = True

In [44]:
N1 = N2 = 75

In [45]:
mesh = dolfin.RectangleMesh(0, 0, 1, 1, N1, N2)

In [46]:
dolfin.RectangleMesh(0, 0, 1, 1, 10, 10)

Out[46]:

### Function space from mesh¶

In [47]:
V = dolfin.FunctionSpace(mesh, 'Lagrange', 1)

DEBUG:FFC:Reusing form from cache.


### Variational problem¶

In [48]:
u = dolfin.TrialFunction(V)

In [49]:
v = dolfin.TestFunction(V)

In [50]:
a = dolfin.inner(dolfin.nabla_grad(u), dolfin.nabla_grad(v)) * dolfin.dx

In [51]:
f1 = dolfin.Constant(1.0)

In [52]:
L1 = f1 * v * dolfin.dx

In [53]:
f2 = dolfin.Expression("x[0]*x[0] + x[1]*x[1]")

In [54]:
L2 = f2 * v * dolfin.dx


### Boundary conditions¶

In [55]:
u0 = dolfin.Constant(0)

In [56]:
def u0_boundary(x, on_boundary):
# try to pin down the function at some interior region:
#if np.sqrt((x[0]-0.5)**2 + (x[1]-0.5)**2) < 0.1:
#    return True
return on_boundary

In [57]:
bc = dolfin.DirichletBC(V, u0, u0_boundary)


### Solve the problem¶

In [58]:
A = dolfin.assemble(a)

DEBUG:FFC:Reusing form from cache.

In [59]:
b = dolfin.assemble(L1)

DEBUG:FFC:Reusing form from cache.

In [60]:
bc.apply(A, b)

In [61]:
u_sol1 = dolfin.Function(V)

In [62]:
dolfin.solve(A, u_sol1.vector(), b)

Out[62]:
1
In [63]:
u_sol2 = dolfin.Function(V)

In [64]:
dolfin.solve(a == L2, u_sol2, bc)

DEBUG:FFC:Reusing form from cache.
DEBUG:FFC:Reusing form from cache.


### Dolfin plot¶

In [65]:
dolfin.plot(u_sol1)
dolfin.interactive()


### Save VTK files¶

In [66]:
dolfin.File('u_sol1.pvd') << u_sol1

In [67]:
dolfin.File('u_sol2.pvd') << u_sol2

In [68]:
f = dolfin.File('combined.pvd')
f << mesh
f << u_sol1
f << u_sol2


### Function evaluation¶

In [69]:
u_sol1([0.21, 0.67])

Out[69]:
0.0466076997781351

### Obtain NumPy arrays¶

In [70]:
u_mat1 = u_sol1.vector().array().reshape(N1+1, N2+1)

In [71]:
u_mat2 = u_sol2.vector().array().reshape(N1+1, N2+1)

In [72]:
X, Y = np.meshgrid(np.linspace(0, 1, N1+2), np.linspace(0, 1, N2+2))

In [73]:
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
cmap = mpl.cm.get_cmap('Reds')

c = ax1.pcolor(X, Y, u_mat1, cmap=cmap)
cb = plt.colorbar(c, ax=ax1)
ax1.set_xlabel(r"$x$", fontsize=18)
ax1.set_ylabel(r"$y$", fontsize=18)
cb.set_label(r"$u(x, y)$", fontsize=18)
cb.set_ticks([0.0, 0.02, 0.04, 0.06])

c = ax2.pcolor(X, Y, u_mat2, cmap=cmap)
cb = plt.colorbar(c, ax=ax2)
ax1.set_xlabel(r"$x$", fontsize=18)
ax1.set_ylabel(r"$y$", fontsize=18)
cb.set_label(r"$u(x, y)$", fontsize=18)
cb.set_ticks([0.0, 0.02, 0.04])

fig.savefig("ch11-fdm-2d-ex1.pdf")
fig.savefig("ch11-fdm-2d-ex1.png")
fig.tight_layout()

In [74]:
X, Y = np.meshgrid(np.linspace(0, 1, N1+1), np.linspace(0, 1, N2+1))

In [75]:
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6), subplot_kw={'projection': '3d'})

p = ax1.plot_surface(X, Y, u_mat1, rstride=4, cstride=4, linewidth=0, cmap=mpl.cm.get_cmap("Reds"))
cb = fig.colorbar(p, ax=ax1, shrink=0.5)
ax1.set_xlabel(r"$x_1$", fontsize=18)
ax1.set_ylabel(r"$x_2$", fontsize=18)
cb.set_label(r"$u(x_1, x_2)$", fontsize=18)

p = ax2.plot_surface(X, Y, u_mat2, rstride=4, cstride=4, linewidth=0, cmap=mpl.cm.get_cmap("Reds"))
cb = fig.colorbar(p, ax=ax2, shrink=0.5)
ax2.set_xlabel(r"$x_1$", fontsize=18)
ax2.set_ylabel(r"$x_2$", fontsize=18)
cb.set_label(r"$u(x_1, x_2)$", fontsize=18)