Solving 2D Heat Equation…

Before we do the Python code, let’s talk about the heat equation and finite-difference method. Heat equation is basically a partial differential equation, it is

If we want to solve it in 2D (Cartesian), we can write the heat equation above like this:

where u is the quantity that we want to know, t is for temporal variable, x and y are for spatial variables, and α is diffusivity constant. So basically we want to find the solution u everywhere in x and y, and over time t.

We can write the heat equation above using finite-difference method like this:

If we arrange the equation above by taking Δx = Δy, we get this final equation:


\(\gamma = \alpha \frac{{\Delta t}}{{\Delta {x^2}}}\)

We use explicit method to get the solution for the heat equation, so it will be numerically stable whenever

\(\Delta t \leq \frac{{\Delta {x^2}}}{{4\alpha }}\)

Everything is ready. Now we can solve the original heat equation approximated by algebraic equation above, which is computer-friendly.

Let’s suppose a thin square plate with the side of 50 unit length. The temperature everywhere inside the plate is originally 0 degree (at t = 0), let’s see the diagram below:

For our model, let’s take Δx = 1 and α = 2.0.
Now we can use Python code to solve this problem numerically to see the temperature everywhere (denoted by i and j) and over time (denoted by k). Let’s first import all of the necessary libraries, and then set up the boundary and initial conditions.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation

print("2D heat equation solver")

plate_length = 50
max_iter_time = 750

alpha = 2
delta_x = 1

delta_t = (delta_x ** 2)/(4 * alpha)
gamma = (alpha * delta_t) / (delta_x ** 2)

# Initialize solution: the grid of u(k, i, j)
u = np.empty((max_iter_time, plate_length, plate_length))

# Initial condition everywhere inside the grid
u_initial = 0

# Boundary conditions
u_top = 100.0
u_left = 0.0
u_bottom = 0.0
u_right = 0.0

# Set the initial condition

# Set the boundary conditions
u[:, (plate_length-1):, :] = u_top
u[:, :, :1] = u_left
u[:, :1, 1:] = u_bottom
u[:, :, (plate_length-1):] = u_right

We’ve set up the initial and boundary conditions, let’s write the calculation function based on finite-difference method that we’ve derived above.

def calculate(u):
    for k in range(0, max_iter_time-1, 1):
        for i in range(1, plate_length-1, delta_x):
            for j in range(1, plate_length-1, delta_x):
                u[k + 1, i, j] = gamma * (u[k][i+1][j] + u[k][i-1][j] + u[k][i][j+1] + u[k][i][j-1] - 4*u[k][i][j]) + u[k][i][j]

    return u

Let’s prepare the plot function so we can visualize the solution (for each k) as a heat map. We use Matplotlib library, it’s easy to use.

def plotheatmap(u_k, k):
    # Clear the current plot figure

    plt.title(f"Temperature at t = {k*delta_t:.3f} unit time")

    # This is to plot u_k (u at time-step k)
    plt.pcolormesh(u_k,, vmin=0, vmax=100)

    return plt

One more thing that we need is to animate the result because we want to see the temperature points inside the plate change over time. So let’s create the function to animate the solution.

def animate(k):
  plotheatmap(u[k], k)

anim = animation.FuncAnimation(plt.figure(), animate, interval=1, frames=max_iter_time, repeat=False)"heat_equation_solution.gif")

That’s it! And here’s the result: