Visualising the Heat Equation

What is the Heat Equation?

#

The term “Heat Equation” refers to a Partial Differential Equation (PDE) of the form: $$ \frac{\partial u}{\partial t} = \kappa \nabla^2 u $$ where \(u = u(\mathbf x, t)\). The parameter \(\kappa\) (sometimes \(\alpha\) is used instead) is the “diffusivity constant”

It arises in many areas of Applied Mathematics and Physics such as

1. Particle diffusion
2. Financial maths
3. Image analysis

For more information, see the Wikipedia page

I wish to solve this to learn more about how we can solve PDEs numerically, and to make a nice visualisation to go along with my numerical solution. First, we will take a look at a way we can solve the equation numerically, and then simply transfer our analytical solution across to the computer. Then we will look at a more “numerical” way of approaching it.

Our Problem

#

The specific heat problem we are going to solve is as follows: $$ \frac{\partial u}{\partial t} = \kappa \frac{\partial^2 u}{\partial x^2} $$ for \(x \in (0,L)\) and \(t \ge 0\) with boundary conditions:

1. \(u(0, t) = 0\)
2. \(u(L, t) = 0\)

and initial condition: \(u(x,0) = f(x)\).

Essentially, we are solving the Heat Equation on a rod with temperature 0 on the boundaries with arbitrary diffusivity constant, rod length and initial condition.

Visualisation

#

I think a sensible place to start is writing the code for visualising a solution. This might be over-engineering, but I like the idea of having an abstract class which implements a visualise method, and forces subclasses to implement some method which gives the solution at a given value at \(x\) and \(t\).

# Abstract class for implementing a solver for the equation.

from abc import ABC, abstractmethod
from typing import Callable
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt
import numpy as np


class Solver(ABC):
    def __init__(self, method_name: str, diffusivity: float, rod_length: float,
                 initial_condition: Callable[[np.ndarray], np.ndarray]):
        self.method_name = method_name

        # These three things well-define our problem
        self.kappa = diffusivity
        self.length = rod_length
        self.initial_condition = initial_condition

    @abstractmethod
    def u(self, x: np.ndarray, t: np.ndarray) -> np.ndarray:
        """ What's the value of the solution at these values of xs and ts?

        Parameters:
            x : np.ndarray
                The values of x to calculate u(x,t) at
            t : np.ndarray
                The values of t to calculate u(x,t) at
        Return Value : np.ndarray
            Shape = (len(t), len(x))
        """

        raise NotImplementedError()

    def visualise(self, xs: np.ndarray, t_max: float, dt: float,
                  out_path=None) -> None:
        """ Visualise the solution

        Parameters:
            xs : np.ndarray
                The values of x to plot the solution at
            t_max: float
                The maximum value of t to plot before looping to t=0
            dt: float
                The difference in time between frames of the animation
        """

        ts = np.arange(dt, t_max + dt, dt)
        us = self.u(xs, ts)

        # Plot these us and animate over t
        fig, ax = plt.subplots()

        fig.suptitle(f"{self.method_name} Solution")

        ax.set_xlabel("x")
        ax.set_ylabel("u(x,t)")
        ax.set_xlim(xs.min(), xs.max())
        ax.set_ylim(us.min(), us.max())

        line, = ax.plot([], [], lw=2)

        def init_anim():
            line.set_data([], [])

        def update_anim(frame):
            line.set_data(xs, us[frame])
            ax.set_title(f"t = {ts[frame]:.2f}")

        anim = FuncAnimation(fig, update_anim, frames=len(ts),
                             init_func=init_anim, blit=False, interval=25,
                             repeat=True)

        if out_path is not None:
            anim.save(out_path, writer="ffmpeg", fps=30)

        plt.show()

Now, we can make a class for each method, implementing u(x,t) as required. For example:

# An example of a Solver implementation
# NOTE: This does NOT solve the heat equation, it's just an example.

from solver import Solver
import numpy as np


class SolverExample(Solver):
    def __init__(self) -> None:
        super().__init__(1, 1, lambda x: 1)

    def u(self, x: np.ndarray, t: np.ndarray) -> np.ndarray:
        return np.sin(x[None, :] - t[:, None]) * np.exp(-0.1 * t[:, None])

Alternatively, if you wanted to do it functionally, you could always just pass the appropriate solving function (in our OOP case: the u(x,t)’s) into the visualise function; but this way was calling to me.

Separation of Variables

#

A popular method for solving PDEs analytically is the Separation of Variables method. We explicitly seek “separable solutions”. In this case, it will be solutions of the form $$ u(x,t) = X(x)T(t) $$

Also, it exploits a key fact about solutions to PDEs:

• Linear Superposition: if the PDE and BCs are linear and homogeneous, the sum of two solutions gives us another solution

The method goes something like this:

1. Ansatz: Assume that \(u = XT\)
2. Substitution: Substitute this ansatz into the original equation, you should uncover ODEs for \(X, T\) coupled by some constant (\(\lambda\)). The ODE for \(X(x)\) is the eigenvalue problem, we solve this first.
3. Solving: Seek non-zero solutions to the eigenvalue problems. We should get eigenvalues \(\lambda_n\) and eigenfunctions \(X_n\) which satisfy the ODE and BCs for appropriate \(n\)
4. Solving: Now solve the ODE for \(T(t)\), we should get more functions \(T_n\).
5. Combining: We can now combine our solutions to get our general solution of the form $$ u(x,t) = \sum_{n=0}^\infty X_n(x) T_n(t) $$
6. Initial Conditions: Now can apply the initial condition for the PDE, this should let us solve for any remaining constants (apart from the ones appearing in the original PDE of course).

For the sake of brevity, I won’t go through the working for this problem. However, the general solution we obtain is: $$ u(x,t) = \sum_{n=0}^\infty A_n \sin \left( \frac{n \pi x}{L} \right) \exp \left( -\left(\frac{n \pi}{L}\right)^2 \kappa t \right) $$ Then we can solve for the \(A_n\) by recognising that \(u(x,0)\) is a (Fourier) sine series. Then we obtain that: $$ A_n = \frac{2}{L} \int_0^L f(x) \sin \left( \frac{n \pi x}{L} \right)\, dx $$

Implementing This Method Numerically

#

Implementing this method is quite simple, especially in this case. We need to do two things:

1. Calculate the \(A_n\) using some numerical integration scheme
2. Truncate the series and sum it like one normally would

My implementation is as follows:

# Solving the heat equation using the series that arises from
# Separation of Variables

from typing import Callable
from solvers.solver import Solver
import numpy as np


class SeriesSolver(Solver):
    def __init__(self, diffusivity: float, rod_length: float,
                 initial_condition: Callable[[np.ndarray], np.ndarray],
                 num_terms: int, n_int_points: int):
        super().__init__("Truncated Series",
                         diffusivity, rod_length, initial_condition)

        # Number of terms that we take from the infinite series
        self.num_terms = num_terms

        # We can actually precompute the coefficients An here since
        # these do not vary per evaluation
        self.coeffs = self.calculate_coefficients(n_int_points)

    def calculate_coefficients(self, n_points: int) -> np.ndarray:
        """ Integrate numerically to find the coefficients of the terms of
        the series (the An)

        """
        # An = 2/L * int_0^L ( f(x) * sin(n pi x / L) ) dx
        coeffs = []

        for n in range(1, self.num_terms + 1):
            coeffs.append(self.integrate(
                lambda x, n=n: np.sin(n * np.pi * x / self.length)
                * self.initial_condition(x),
                0,
                self.length,
                n_points
            ))

        return (2 / self.length) * np.array(coeffs)

    def integrate(self, f: Callable[[np.ndarray], np.ndarray],
                  lower_limit: float, upper_limit: float,
                  n_points: int) -> float:
        """ Integrate f(x) between the limits (inclusive) using (Composite)
        Simpson's Rule

        Parameters:
            f : Callable[[np.ndarray], np.ndarray]
                The integrand
            lower_limit : float
            upper_limit : float
            n_points : int
                The number of points to use in the approximation,
                if it's not odd, it will be rounded up to the next
                odd number
        """

        # Ensure n_points is positive and odd
        n_points = abs(n_points)
        if n_points < 2:
            n_points = 2
        if n_points % 2 == 0:
            n_points += 1

        xs = np.linspace(lower_limit, upper_limit, num=n_points)
        ys = f(xs)

        h = xs[1] - xs[0]

        s = ys[0] + ys[-1] + 4 * ys[1:-1:2].sum() + 2 * ys[2:-1:2].sum()
        return h * s / 3

    def u(self, x: np.ndarray, t: np.ndarray) -> np.ndarray:
        x = np.asarray(x)
        t = np.asarray(t)

        n = np.arange(1, self.num_terms + 1)

        spatial_part = np.sin(np.outer(n, np.pi * x / self.length))
        decay_part = np.exp(-self.kappa * ((n * np.pi / self.length) ** 2)
                            * t[:, None])

        weighted_decay = decay_part * self.coeffs[None, :]

        return weighted_decay @ spatial_part

To see what the solution looks like, let’s visualise our problem with:

• \(\kappa = 0.1\)
• \(L = 1\)
• \(f(x) = \sin(2 \pi x)\)

This visualisation was created using the series_visualisation script available in the repository.

Finite Difference Methods

#

Another way we can approach solving PDEs numerically is using finite difference schemes. In short, we discretise the domain and solve the differential equation on a grid of points. Specifically, this method is known as the Method of Lines

To avoid going into too much detail in this post, I will simply give a brief outline of how we are going to apply it to the Heat Equation.

1. Discretise in Space: Split the spatial domain \((0, L)\) into points \(x_1, x_2, \dots, x_n\). From this, we can uncover coupled differential equations for \(u_j(t) = u(x_j, t)\)
2. Organise: We then organise these differential equations into a vector ODE.
3. Discretise in Time: Finally we can discretise the resulting ODE with respect to \(t\), by considering \(\mathbf u^j = \mathbf{u}(t_j) = (u(x_1, t_j), \dots, u(x_n, t_j))\).

Let \(\Delta t, \Delta x\) be the spacing between subsequent grid points. For space derivatives, we will use the central difference approximation: $$ \frac{\partial^2 u}{\partial x^2}(x=x_j)= \frac{u_{j-1} - 2u_j + u_{j+1}}{\Delta x^2} $$

For time derivatives, we will use the trapezoidal scheme: $$ \frac{\mathbf u^{j+1} - \mathbf u^j}{\Delta t} = \frac{A \mathbf u^{j+1} + A \mathbf u^j}{2} $$ where \(A\) is the matrix we uncover when organising the space ODEs into a single vector ODE.

Using the method of lines, with these finite difference schemes on the Heat Equation is also known as the Crank-Nicolson Scheme. When we follow this method through, we get that: $$ \left(I - \frac{\Delta t}{2} A\right) \mathbf u^{j+1} = \left(I + \frac{\Delta t}{2} A\right) \mathbf u^j $$ where $$ A = \frac{1}{\Delta x^2} \begin{pmatrix} -2 & 1 \\ 1 & -2 & 1 \\ & 1 & -2 & 1 \\ & & \ddots & \ddots & \ddots \\ & & & 1 & -2 & 1 \\ & & & & 1 & -2 & 1 \\ & & & & & 1 & -2 \end{pmatrix} $$

Implementing This Method

#

Unlike the series solution, we need to be a bit more considerate about how we implement it since the naive method would be to take the inverse of matrix sum on the left. However, inverses are computationally expensive.

Instead, we take the LU decomposition and solve the system of equations in that way. This is much cheaper!

My implementation is as follows:

# Solving the heat equation using Finite Difference Schemes
# i.e. the Method of Lines, specifically the Crank-Nicholson Scheme

from typing import Callable
from solvers.solver import Solver
from scipy.linalg import lu_factor, lu_solve
import numpy as np


class MoLSolver(Solver):
    def __init__(self, diffusivity: float, rod_length: float,
                 initial_condition: Callable[[np.ndarray], np.ndarray]):
        super().__init__("Method of Lines",
                         diffusivity, rod_length, initial_condition)

    def u(self, x: np.ndarray, t: np.ndarray) -> np.ndarray:
        # (I - dt/2 A) * u_n+1 = (I + dt/2 A) * u_n

        # We're going to implement this with the assumption that the passed
        # array of x and t values ARE uniformly spaced points.
        dx = x[1] - x[0]
        dt = t[1] - t[0]

        # The number of interior points on the x grid
        m = len(x) - 2

        # The recurrence essentially comes down to solving a system of
        # equations; we COULD take the inverse but this is very expensive
        # so we use LU decomposition instead.

        # This requires a matrix constructed in a specific way
        r = self.kappa * dt / (dx ** 2)
        tridiag = np.diag(-2 * np.ones(m)) \
            + np.diag(np.ones(m - 1), 1) \
            + np.diag(np.ones(m - 1), -1)

        id = np.eye(m)
        lhs = id - 0.5 * r * tridiag
        rhs = id + 0.5 * r * tridiag

        lu, piv = lu_factor(lhs)

        # Initialise solution array with the initial condition
        sol = np.zeros((len(t), len(x)))
        sol[0, :] = self.initial_condition(x)

        # Timestep the vector ODE
        current_u = sol[0, 1:-1]
        for n in range(1, len(t)):
            current_u = lu_solve((lu, piv), rhs @ current_u)
            sol[n, 1:-1] = current_u

        # Boundary conditions = 0; so we don't need to do anything at the
        # boundaries

        return sol

If we visualise the problem with the same parameters as with the series solver, we obtain a very familiar animation:

This visualisation was created with the mol_visualisation script available in the repository.

Comparing the Methods

#

So why do we have different methods for doing the same thing? One must be definitively better than the other… right?

Well let’s consider the pros of the series solution first:

1. The full analytical solution is exact, so the only error we obtain in the implementation is the truncation and integration error.
2. The implementation is much more straightforward and each term of the series, no matter how complex it is, can be computed in parallel.

However, it has some drawbacks:

1. It was easy to do in this case since we had a simple domain: a rod. The difficulty ramps up when we consider more complex geometries and PDEs.
2. The coefficients are obtained by integration which can introduce more error into the solution.

As for the Crank-Nicolson scheme:

1. We can generalise it to more complex geometries: we just need to set up the grid appropriately.
2. This scheme is stable for any choice of timestep, while still being decently simple to implement.

Nonetheless:

1. There’s more room for error since we are discretising a continuous domain
2. If we have a very fine grid (i.e. many grid points), we’d have to solve a very large system of linear equations.
3. Accuracy is dependent on choosing a small enough step size.

In short: the series solution is elegantly straightforward for problems with simple boundary conditions and domains, but doesn’t scale well to complex PDEs. For these, it is easier to turn to finite difference methods, like the Method of Lines, which give us a general recipe to tackle any PDE: discretise, approximate, solve.

Therefore, it’s quite easy to see how analytical solutions are highly elegant and can reveal the structure of the solution; whereas numerical methods are built for flexibility to allow us to solve a large range of PDEs, even those we cannot solve analytically.

It is important to note, however, that there are many more methods to solve PDEs which we haven’t considered, e.g. Spectral Methods which borrows key ideas from the Separation of Variables method.

Repository

#

The repository containing all the code included in this article, and the scripts used to generate the graphs can be found here:

jenboc/heat-equation-visualisation

Python

0

0