Python Solution of the Diffusion Equation | Lecture 73 | Numerical Methods for Engineers

diffusion_equation_solution.py

import matplotlib.pyplot as plt
import numpy as np
import scipy.sparse
import scipy.sparse.linalg


def main() -> None:
    diffusion = 1
    length = 1
    n_points = 500
    n_out = 500
    n_steps = 10_000

    x = np.linspace(-length, length, n_points)
    dx = x[1] - x[0]
    dt = dx**2 / (2 * diffusion)
    alpha = dt * diffusion / (dx**2)
    m = scipy.sparse.diags(
        [-alpha, 2 * (1 + alpha), -alpha], [-1, 0, 1], shape=(n_points, n_points)
    )
    m = m.tolil()
    m[[0, -1], :] = 0
    m[[0, -1], [0, -1]] = 1
    m = m.tocsr()

    sigma = length / 16
    u = 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-0.5 * (x / sigma) ** 2)

    _, ax = plt.subplots(constrained_layout=True)
    ax.grid(True)

    ax.plot(x, u)
    for idx in range(1, n_steps):
        u[1:-1] = alpha * u[:-2] + 2 * (1 - alpha) * u[1:-1] + alpha * u[2:]
        u = scipy.sparse.linalg.spsolve(m, u)
        if not np.mod(idx, n_out):
            ax.plot(x, u)


if __name__ == "__main__":
    main()
    plt.show()