Solving PDEs with the Fourier spectral method in 2D#
We will discuss the Fourier spectral method for solving PDEs and focus on the 2D Poisson equation and the heat equation.
Fourier techniques in 2D#
As in the one-dimensional, we can perform a Fourier expansion of the solution \(u(x, y)\), where the Fourier coefficients in 2D are simply given by computing 1D Fourier coefficients in each direction.
with \(\mathbf{x} = (x, y)\) and \(\mathbf{k} = (k_x/L_x, k_y/L_y)\).
Being a bit in rush :), we simply summarize here that one can basically develop many of the Fourier techniques we discussed for the 1D case in the 2D case, just simply applying all the relevant concepts to each spatial direction separately.
For example, the formal 2D Fourier series of a function \(f(x, y)\) is given by
where the 2D Fourier coefficients are defined as above.
As in the 1D case, we can then derive the following identities for derivatives of periodic functions and their corresponding Fourier coefficients:
Due to the appearance of the factors \(2\pi/L_x\) and \(2\pi/L_x\) it also very common to ease the notation by defining the wavenumber vector \(\mathbf{\tilde{k}}\) as
This way, the Laplacian in Fourier space becomes simply
Let’s use this now to solve the Poisson equation in 2D.
2D Poisson equation#
Let’s consider the 2D Poisson equation is given by
on a domain \(\Omega = [0, L_x) \times [0, L_y)\) supplemented with periodic boundary conditions.
To solve this equation on a continuous level, we can do a Fourier expansion of the solution \(u(x, y)\) and the right-hand side \(f(x, y)\), and then solve for the Fourier coefficients, exactly as you did in the 1D case back in the Matte 4K course. Then the Fourier coefficients of the solution are given by
Note that this can only be done if the wavenumber vector \(\mathbf{\tilde{k}}\) is not zero, i.e., \(\mathbf{\tilde{k}} \neq 0\). But observe that the Fourier coefficients for \((k_x, k_y) = (0, 0)\)
is simply the mean value of the solution \(u(x, y)\) over the domain \(\Omega\). The division by \(0\) “problem” is directly related to the fact that there is a ambiguity in the solution of the Poisson equation, since the Laplacian of a constant (and thus periodic!) function is zero, and therefore for any solution \(u\) the function \(u + c\) is also a solution, and thus the solution is only determined up to a constant. To eliminate this ambiguity, the convention is to prescribe the mean value of the solution, for instance to zero, i.e., by requiring \(\int_{\Omega} u(x, y) \;\mathrm{d}x \;\mathrm{d}y = 0\). This will then uniquely determine the zero mode \(\widehat{u}(\mathbf{0})\).
Note that this is not a problem if you e.g. want to solve the Poisson problem with a lower order term, i.e., if you have a Poisson equation of the form
for some constant \(c>0\), since then the solution on the Fourier side is given by
This make sense, because in this case, adding a non-zero constant function to \(u\) will change the right-hand side of the Poisson equation and thus the solution is unique.
We will now exploit these formulas and ideas numerically, and solve the Poisson equation numerically by using the (2 dimensional) fast fourier transform (FFT) to approximate the Fourier coefficients of the right-hand side \(f(x, y)\), divide then by the norm of the wavenumber vector, and then use the inverse FFT to compute the solution \(u(x, y)\). This is the so-called Fourier spectral method. Le’t see how this works in the following code.
# %matplotlib widget
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fft, fft2, ifft, ifft2, fftfreq, fftshift
import pandas as pd
First, we consider a periodic function \(u(x, y)\) on the domain \(\Omega = [0, 2\pi) \times [0, 2\pi)\) for which we can easily compute the corresponding right-hand side \(f(x, y)\).
Let’s use the following function
We can easily compute the right-hand side \(f(x, y)\) by inserting this function into the Poisson equation and obtaining
We could also say the the given \(u\) is an eigenfunction of the Laplacian with eigenvalue \(-5\). Let’s start by plotting the function \(u(x, y)\) and the right-hand side \(f(x, y)\).
Here, we need some constructs from the numpy and matplotlib libraries.
# Define the domain
Lx, Ly = 2*np.pi, 2*np.pi
# Define 1d samplings for x and y directions
# Nx, Ny = 32, 32
Nx, Ny = 3, 4
x = np.linspace(-Lx/2, Lx/2, Nx, endpoint=False)
y = np.linspace(-Ly/2, Ly/2, Ny, endpoint=False)
# Generate a 2d sampling grid to be evaluate functions of x and y
X, Y = np.meshgrid(x, y)
# X, Y = np.meshgrid(x, y, sparse=True)
print(f"X = {X}")
print(f"Y = {Y}")
X = [[-3.14159265 -1.04719755 1.04719755]
[-3.14159265 -1.04719755 1.04719755]
[-3.14159265 -1.04719755 1.04719755]
[-3.14159265 -1.04719755 1.04719755]]
Y = [[-3.14159265 -3.14159265 -3.14159265]
[-1.57079633 -1.57079633 -1.57079633]
[ 0. 0. 0. ]
[ 1.57079633 1.57079633 1.57079633]]
# We can also define a sparse 2d sampling grid
X, Y = np.meshgrid(x, y, sparse=True)
print(f"X = {X}")
print(f"Y = {Y}")
X = [[-3.14159265 -1.04719755 1.04719755]]
Y = [[-3.14159265]
[-1.57079633]
[ 0. ]
[ 1.57079633]]
Now that we have understood how the meshgrid arrays look like, let’s use a finer mesh.
Nx, Ny = 32, 32
x = np.linspace(-Lx/2, Lx/2, Nx, endpoint=False)
y = np.linspace(-Ly/2, Ly/2, Ny, endpoint=False)
# Generate a 2d sampling grid to be evaluate functions of x and y
X, Y = np.meshgrid(x, y, sparse=True)
# Define the solution
def u_ex(x, y):
return np.sin(x)*np.cos(2*y)
# Evaluate the exact solution on the grid
U_ex = u_ex(X, Y)
# Plot the exact solution as surface plot
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(X, Y, U_ex, cmap='viridis', antialiased=True)
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$U_\mathrm{ex}$')
plt.show()
Let’a also plot the right-hand side \(f(x, y)\).
# Define the solution
def f(x, y):
return 5*np.sin(x)*np.cos(2*y)
# Evaluate the exact solution on the grid
F = f(X, Y)
# Plot the exact solution as surface plot
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(X, Y, F, cmap='viridis', antialiased=True)
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$F$')
plt.show()
Now we write a tiny function consisting of basically 10 lines of which computes the solution \(u(x, y)\) on the Fourier side by using the formula
def solve_poisson(F, Lx, Ly, Nx, Ny):
# Compute the FFT of the right-hand side using the 2D FFT
F_hat = fft2(F)
# Compute wave number grid
kx = fftfreq(Nx, d=Lx/Nx)*2*np.pi
ky = fftfreq(Ny, d=Ly/Ny)*2*np.pi
KX, KY = np.meshgrid(kx, ky, sparse=True)
# Compute the Poisson operator in Fourier space
K2 = KX**2 + KY**2
# Just modified to avoid division by zero
# We set the zero frequency component to 0 explicitly below
K2[0, 0] = 1
# Solve the Poisson equation in Fourier space
U_hat = F_hat / K2
# Set the zero frequency component to zero
# This corresponds to setting the average value of the solution to zero
U_hat[0, 0] = 0
# Compute the inverse 2D FFT to get the solution
return U_hat
Let’s apply this function and plot the exact solution \(U_{\mathrm{ex}}(x, y)\), the numerical solution and the error \(U - U_{\mathrm{ex}}\).
U_hat = solve_poisson(F, Lx, Ly, Nx, Ny)
U = ifft2(U_hat).real
# Plot the solution
fig = plt.figure(figsize=(15, 5))
ax = fig.add_subplot(131, projection='3d')
surf = ax.plot_surface(X, Y, U_ex, cmap='viridis', antialiased=True)
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$U_\mathrm{ex}$')
ax = fig.add_subplot(132, projection='3d')
surf = ax.plot_surface(X, Y, U, cmap='viridis')
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$U(X, Y)$')
U_err = U - U_ex
print(f"Error norm: {np.abs(U_err).max()}")
ax = fig.add_subplot(133, projection='3d')
surf = ax.plot_surface(X, Y, U_err, cmap='viridis')
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$U-U_{\mathrm{ex}}$')
plt.show()
Error norm: 6.661338147750939e-16
Let’s try a more complicated manufactured solution. Since computing the rhs is tedious and error-prone,
we will use the sympy library to compute the rhs. We will
pass the exact solution as string, the following function will then compute the rhs for us,
and return the exact solution and the rhs as numpy compatible functions.
# Exact solution and its -Laplacian
def manufacture_solution_poisson(u_str):
"""
Generate the exact solution and its corresponding right-hand side for the Poisson equation.
This function takes a string representation of the exact solution `u(x, y)` and computes
its Laplacian to generate the corresponding right-hand side `f(x, y)` for the Poisson equation.
The function returns `u(x, y)` and `f(x, y)` as `numpy`-compatible callable functions.
Parameters:
u_str (str): A string representation of the exact solution `u(x, y)`.
Returns:
tuple: A tuple containing two functions:
- u (function): The exact solution `u(x, y)` as a `numpy`-compatible function.
- f (function): The right-hand side `f(x, y)` as a `numpy`-compatible function.
"""
import sympy as sy
from sympy import sin, cos, exp
x, y = sy.symbols('x y')
u_sy = eval(u_str)
laplace = lambda u: sy.diff(u, x, x) + sy.diff(u, y, y)
f_sy = -sy.simplify(laplace(u_sy))
print(f'u = {u_sy}')
print(f'f = {f_sy}')
u = sy.lambdify((x, y), u_str, modules='numpy')
f = sy.lambdify((x, y), f_sy, modules='numpy')
return u, f
# u_ex_str = 'sin(x)*cos(2*y)'
u_ex_str = 'exp(sin(x)) + cos(2*y)'
u_ex, f = manufacture_solution_poisson(u_ex_str)
u = exp(sin(x)) + cos(2*y)
f = -(-sin(x) + cos(x)**2)*exp(sin(x)) + 4*cos(2*y)
Let’s solve the Poisson equation with the new manufactured solution.
# Example usage
Lx, Ly = 2*np.pi, 2*np.pi
# Nx, Ny = 10,10,
Nx, Ny = 32, 32
x = np.linspace(-Lx/2, Lx/2, Nx, endpoint=False)
y = np.linspace(-Ly/2, Ly/2, Ny, endpoint=False)
# Generate a 2d sampling grid to be evaluate functions of x and y
X, Y = np.meshgrid(x, y, sparse=True)
F = f(X, Y)
U_hat = solve_poisson(F, Lx, Ly, Nx, Ny)
U = ifft2(U_hat).real
# Re-adjust solution to same mean value as exact solution
# Only necessary for comparison with manufactured solution
# which does not have zero mean
U_ex = u_ex(X, Y)
U += np.mean(U_ex)
U_err = U - U_ex
err = np.abs(U_err).max()
print(f'Error: {err}')
Error: 1.7763568394002505e-15
And let’s plot the exact solution, the numerical solution and the error.
# Plot the solution
fig = plt.figure(figsize=(15, 5))
ax = fig.add_subplot(131, projection='3d')
surf = ax.plot_surface(X, Y, U_ex, cmap='viridis', antialiased=True)
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$U_\mathrm{ex}$')
ax = fig.add_subplot(132, projection='3d')
surf = ax.plot_surface(X, Y, U, cmap='viridis')
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$U(X, Y)$')
ax = fig.add_subplot(133, projection='3d')
surf = ax.plot_surface(X, Y, U_err, cmap='viridis')
fig.colorbar(surf, shrink=0.6)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel(r'$U-U_{\mathrm{ex}}$')
plt.show()
A last thing: Since the 3D surface plots are rather slow, there are not really suitable when visualizing solutions on fine grids or a lot of snapshots.
You can use the imshow_plot_u function from our homecooked little wrapper module project_tools to plot the solution as a 2D image:
import os
import os.path
import sys
# Add path to project_tools.py to Python's search path
project_tools_path = os.path.join(os.getcwd(), '../project_3_2025')
if project_tools_path not in sys.path:
sys.path.append(project_tools_path)
import project_tools as pot
pot.imshow_plot_u(U, Lx, Ly, cblabel=r'$U(X, Y)$')
(<Figure size 640x480 with 2 Axes>, <Axes: xlabel='$x$', ylabel='$y$'>)
