import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as mplanim
from matplotlib import (
rc,
) # Configure matplotlib to render animations in Jupyter notebooks
rc("animation", html="html5")
Solving a PDE Numerically in Python¶
The Heat Equation¶
The heat equation is a PDE that models the diffusion of heat through a region. In this first example, we will solve a transient heat equation for a 2D domain, $\Omega = (0, 10) \times (0, 10)$ and $t > 0$. It is given as: \begin{equation} \tag{1} \frac{\partial u}{\partial t} - \alpha (\frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}) = 0 \quad x \in [0, 10] \quad y \in [0, 10] \quad t \in [0, 1] \end{equation} where $u [K]$ is the temperature, $x, y [m]$ are the spatial coordinates, $t [s]$ is time, and $\alpha$ is the thermal diffusivity of the domain, which we choose to be $\alpha = 0.4$.
It has Dirichlet BCs on the bottom, top and left sides given as: $$ u(x, 0, t) = u(x, 1, t) = u(0, y, t) = 0 K $$
With the right side given as: $$ u(1, y, t) = 100 K $$
It has an IC inside throughout the domain given by: $$ u(x, y, 0) = 0 K \quad x \in [0, 10] \quad y \in [0, 10] $$
Let's start by creating our domain $\Omega$ for the PDE $u$ with our BCs and IC assigned.
# Domain size
x_len = 100
y_len = 100
# Number of time steps
max_time = 25
# Thermal diffusivity
alpha = 0.4
# Spatial mesh sizes && Number of spatial mesh nodes
delta_x = 1
delta_y = 1
nx, ny = int(x_len / delta_x), int(y_len / delta_y)
The BCs and ICs for the defined problem are now assigned.
# Initialize solution: u(j, i, n)
u = np.zeros((nx, ny, max_time))
# Initial condition
u_initial = 0.0
# Boundary conditions (Dirichlet)
u_right = 100.0
u_left = 0.0
u_bottom = 0.0
u_top = 0.0
# Set the initial condition
u.fill(u_initial)
# Set the boundary conditions
u[:, (nx - 1) :, :] = u_right
u[:, :1, :] = u_left
u[:1, 1 : (nx - 1), :] = u_bottom
u[(ny - 1) :, 1 : (nx - 1), :] = u_top
# Visualize domain at t=0
fig, ax = plt.subplots(layout="constrained")
ax.set_title(rf"$u(x, y, {0})$, at $t={0}$s")
hm = ax.pcolormesh(
u[:, :, 0],
shading="auto",
cmap=plt.cm.jet,
vmin=u[:, :, 0].min(),
vmax=u[:, :, 0].max(),
)
fig.colorbar(mappable=hm, ax=ax)
Finite Differences in Python¶
We will solve this example using finite differences in Python, with mesh sizes of $\Delta x = 0.1, \Delta y = 0.1$. To ensure stability of our method, the mesh size in $t$ will be determined using the Von Neumann stability criteria. For the 2D heat equation, this criteria is given as: \begin{equation} \tag{2} \Delta t = \frac{\Delta x^2 \cdot \Delta y^2}{2 \alpha (\Delta x^2 +\Delta y^2)} \end{equation}
# Stability calcs
delta_t = min((delta_x**2 * delta_y**2) / (2 * alpha * (delta_x**2 + delta_y**2)), 0.5)
We then discretize equation 1 using the mesh sizes, $\Delta x, \Delta y, \Delta t$. The first derivative in time, $\frac{\partial u}{\partial t}$, is approximated using FD while the second derivatives in space, $(\frac{\partial^2 u}{\partial x^2}, \frac{\partial^2 u}{\partial y^2})$ are approximated using CDs: \begin{equation} \tag{3} \frac{u_{i, j}^{n + 1} - u_{i, j}^{n}}{\Delta t} - \alpha (\frac{u_{i+1, j}^{n} - 2 \cdot u_{i, j}^{n} + u_{i-1, j}^{n}}{\Delta x^2} + \frac{u_{i, j+1}^{n} - 2 \cdot u_{i, j}^{n} + u_{i, j-1}^{n}}{\Delta y^2}) = 0 \end{equation}
This can then be further simplified to determine the diffusion of heat throughout our domain: \begin{equation} \tag{4} u_{i, j}^{n + 1} = u_{i, j}^{n} + \alpha \cdot \Delta t (\frac{u_{i+1, j}^{n} - 2 \cdot u_{i, j}^{n} + u_{i-1, j}^{n}}{\Delta x^2} + \frac{u_{i, j+1}^{n} - 2 \cdot u_{i, j}^{n} + u_{i, j-1}^{n}}{\Delta y^2}) \end{equation}
This can then be written in Python as:
def calc_u_fd(u, use_opt=True):
"""Finite difference Method for 2D Heat equation
Args:
u (array): Initial temperature distribution at nodes (j, i, n)
use_opt (bool): Use optimised finite differences (only a time loop)
Returns:
array: Final temperature distribution
"""
if not use_opt:
for n in range(0, max_time - 1, 1):
for i in range(1, x_len - 1, delta_x):
for j in range(1, y_len - 1, delta_y):
u_xx = (u[j, i + 1, n] - 2 * u[j, i, n] + u[j, i - 1, n]) / (
delta_x**2
)
u_yy = (u[j + 1, i, n] - 2 * u[j, i, n] + u[j - 1, i, n]) / (
delta_y**2
)
u[j, i, n + 1] = u[j, i, n] + (alpha * delta_t) * (u_xx + u_yy)
else:
for n in range(0, max_time - 1, 1):
u0 = u.copy()
u_xx = (u0[2:, 1:-1, n] - 2 * u0[1:-1, 1:-1, n] + u0[:-2, 1:-1, n]) / (
delta_x**2
)
u_yy = (u0[1:-1, 2:, n] - 2 * u0[1:-1, 1:-1, n] + u0[1:-1, :-2, n]) / (
delta_y**2
)
u[1:-1, 1:-1, n + 1] = u0[1:-1, 1:-1, n] + (alpha * delta_t) * (u_xx + u_yy)
return u
u = calc_u_fd(u, use_opt=False)
def update_u(n, u):
"""Animation showing heat diffusion throughout domain
Args:
n (int): Current time step
u (array): Final temperature distribution
"""
ax.set_title(rf"$u(x, y, {n})$, at $t={n}$s")
hm = ax.pcolormesh(u[:, :, n], shading="auto", cmap=plt.cm.jet)
return hm
ani = mplanim.FuncAnimation(
fig=fig, func=update_u, fargs=(u,), interval=150, frames=max_time, repeat=True
)
ani
Practice¶
Now that we have the code, we can implement something more interesting. Here are some questions that you can use to explore the heat distributions for different conditions.
Example 1:¶
Set left and right BCs to $0.0$ K and the top and bottom to $100.0$ K. The interior has ICs set to $0.0$ K.
Let's first initialize our solution array, $u$:
# Initialize solution: u(j, i, n)
u = np.zeros((nx, ny, max_time))
Next, set the BCs and ICs as described above:
# Initial condition
u_initial = ""
# Boundary conditions (Dirichlet)
u_right = ""
u_left = ""
u_bottom = ""
u_top = ""
# Set the initial condition
u[:, :, 0] = ""
Finally, we will assign the BCs to the sides of the domain.
# Set the boundary conditions
u[:, (nx - 1) :, :] = u_right
u[:, :1, :] = u_left
u[:1, 1 : (nx - 1), :] = u_bottom
u[(ny - 1) :, 1 : (nx - 1), :] = u_top
# Visualize domain at t=0
fig, ax = plt.subplots(layout="constrained")
ax.set_title(rf"$u(x, y, {0})$, at $t={0}$s")
hm = ax.pcolormesh(
u[:, :, 0],
shading="auto",
cmap=plt.cm.jet,
vmin=u[:, :, 0].min(),
vmax=u[:, :, 0].max(),
)
fig.colorbar(mappable=hm, ax=ax)
Now let's solve our system as before.
u = ""
ani = mplanim.FuncAnimation(
fig=fig, func=update_u, fargs=(u,), interval=150, frames=max_time, repeat=True
)
ani
Example 2:¶
Set top, bottom, left and right BCs to $0.0$ K and the interior with ICs set to $0.0$ K except for a circular region with a temperature of $50.0$ K. The circular region has a radius, $r$, of 2 centered at $(5,5)$ within the domain.
Now let's initialize our solution array, $u$, and assign the BCs and ICs.
# Initialize solution: u(j, i, n)
u = ""
# Initial condition
u_initial = ""
# Boundary conditions (Dirichlet)
u_right = ""
u_left = ""
u_bottom = ""
u_top = ""
# Set the initial condition
u[:, :, 0] = ""
The inner circle is assigned using the code below.
# Set initial conditions for inner circle
r, cx, cy = 20, 50, 50
for i in range(nx):
for j in range(ny):
# Calculate distance of point from centre
p = (i * delta_x - cx) ** 2 + (j * delta_y - cy) ** 2
# Assign if point < r
if p < r**2:
u[i, j, 0] = 50.0
We can then visualize our system before solving.
# Visualize domain at t=0
fig, ax = plt.subplots(layout="constrained")
ax.set_title(rf"$u(x, y, {0})$, at $t={0}$s")
hm = ax.pcolormesh(
u[:, :, 0],
shading="auto",
cmap=plt.cm.jet,
vmin=u[:, :, 0].min(),
vmax=u[:, :, 0].max(),
)
fig.colorbar(mappable=hm, ax=ax)
Finally, let's solve this system.
u = ""
ani = mplanim.FuncAnimation(
fig=fig, func=update_u, fargs=(u,), interval=150, frames=max_time, repeat=True
)
ani
References¶
- The Heat Equation/12%3A_Fourier_Solutions_of_Partial_Differential_Equations/12.01%3A_The_Heat_Equation#title)
- Finite difference method for 2-D heat equation
- Stability Analysis in a Transient Problem