Solver Comparison: Regular vs Staggered Grids¶

spectraldiffx provides spectral Poisson/Helmholtz solvers for four distinct grid–BC combinations, plus an FFT solver for periodic domains. This notebook compares all five on a single canonical problem to build intuition for when to use each.

Solver	Transform	BC type	Grid type	Eigenvalue denominator
`solve_poisson_dst`	DST-I	Dirichlet	Regular (vertex-centred)	$2(N+1)$
`solve_poisson_dst2`	DST-II	Dirichlet	Staggered (cell-centred)	$2N$
`solve_poisson_dct1`	DCT-I	Neumann	Regular (vertex-centred)	$2(N-1)$
`solve_poisson_dct`	DCT-II	Neumann	Staggered (cell-centred)	$2N$
`solve_poisson_fft`	FFT	Periodic	—	$N$

Key insight: "regular" means grid points sit on the boundary (vertices); "staggered" means grid points sit at cell centres, with boundaries located half a grid spacing outside the first/last point.

In [ ]:

Copied!





from pathlib import Path

import jax
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

matplotlib.use("Agg")
jax.config.update("jax_enable_x64", True)

from spectraldiffx import (
    dct1_eigenvalues,
    dct2_eigenvalues,
    dst1_eigenvalues,
    dst2_eigenvalues,
    fft_eigenvalues,
    solve_poisson_dct,
    solve_poisson_dct1,
    solve_poisson_dst,
    solve_poisson_dst2,
    solve_poisson_fft,
)

IMG_DIR = Path(__file__).resolve().parent.parent / "docs" / "images" / "solver_comparison"
IMG_DIR.mkdir(parents=True, exist_ok=True)
from pathlib import Path

import jax
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

matplotlib.use("Agg")
jax.config.update("jax_enable_x64", True)

from spectraldiffx import (
    dct1_eigenvalues,
    dct2_eigenvalues,
    dst1_eigenvalues,
    dst2_eigenvalues,
    fft_eigenvalues,
    solve_poisson_dct,
    solve_poisson_dct1,
    solve_poisson_dst,
    solve_poisson_dst2,
    solve_poisson_fft,
)

IMG_DIR = Path(__file__).resolve().parent.parent / "docs" / "images" / "solver_comparison"
IMG_DIR.mkdir(parents=True, exist_ok=True)

1. Grid Geometry¶

The fundamental difference between regular and staggered grids is where the grid points sit relative to the domain boundary.

Regular (vertex-centred)            Staggered (cell-centred)

  boundary                            boundary (half-spacing outside)
  |                                   |
  o----o----o----o----o               |--x----x----x----x--|
  |    |    |    |    |               |  |    |    |    |  |
  o    o    o    o    o               |--x    x    x    x--|
  |    |    |    |    |               |  |    |    |    |  |
  o    o    o    o    o               |--x    x    x    x--|
  |    |    |    |    |               |  |    |    |    |  |
  o----o----o----o----o               |--x----x----x----x--|
  |                                   |
  boundary                            boundary (half-spacing outside)

  o = vertex (on boundary)            x = cell centre (interior)
  N=5 spans [0, L]                    N=4 spans [dx/2, L-dx/2]
  dx = L/(N-1)                        dx = L/N

This distinction matters because the spectral transform that diagonalises the discrete Laplacian depends on where the unknowns live.

In [ ]:

Copied!





# Build small grids for visualisation
N_vis = 6
L = 1.0

# Regular grid: points at vertices, including boundary
x_reg = np.linspace(0, L, N_vis)
y_reg = np.linspace(0, L, N_vis)

# Staggered grid: points at cell centres
dx_stag = L / N_vis
x_stag = np.linspace(dx_stag / 2, L - dx_stag / 2, N_vis)
y_stag = np.linspace(dx_stag / 2, L - dx_stag / 2, N_vis)

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# Regular grid
ax = axes[0]
Xr, Yr = np.meshgrid(x_reg, y_reg)
ax.plot(Xr.ravel(), Yr.ravel(), "ko", ms=8, label="Grid points (vertices)")
ax.axhline(0, color="C3", ls="--", lw=2, label="Boundary")
ax.axhline(L, color="C3", ls="--", lw=2)
ax.axvline(0, color="C3", ls="--", lw=2)
ax.axvline(L, color="C3", ls="--", lw=2)
ax.set_title("Regular (vertex-centred)", fontsize=13)
ax.set_xlim(-0.15, L + 0.15)
ax.set_ylim(-0.15, L + 0.15)
ax.set_aspect("equal")
ax.legend(loc="upper right", fontsize=9)

# Staggered grid
ax = axes[1]
Xs, Ys = np.meshgrid(x_stag, y_stag)
ax.plot(Xs.ravel(), Ys.ravel(), "s", color="C0", ms=8, label="Grid points (cell centres)")
ax.axhline(0, color="C3", ls="--", lw=2, label="Boundary (off-grid)")
ax.axhline(L, color="C3", ls="--", lw=2)
ax.axvline(0, color="C3", ls="--", lw=2)
ax.axvline(L, color="C3", ls="--", lw=2)
ax.set_title("Staggered (cell-centred)", fontsize=13)
ax.set_xlim(-0.15, L + 0.15)
ax.set_ylim(-0.15, L + 0.15)
ax.set_aspect("equal")
ax.legend(loc="upper right", fontsize=9)

plt.suptitle("Grid Geometry: Regular vs Staggered", fontsize=14, y=1.02)
plt.tight_layout()
fig.savefig(IMG_DIR / "grid_geometry.png", dpi=150, bbox_inches="tight")
plt.show()
# Build small grids for visualisation
N_vis = 6
L = 1.0

# Regular grid: points at vertices, including boundary
x_reg = np.linspace(0, L, N_vis)
y_reg = np.linspace(0, L, N_vis)

# Staggered grid: points at cell centres
dx_stag = L / N_vis
x_stag = np.linspace(dx_stag / 2, L - dx_stag / 2, N_vis)
y_stag = np.linspace(dx_stag / 2, L - dx_stag / 2, N_vis)

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# Regular grid
ax = axes[0]
Xr, Yr = np.meshgrid(x_reg, y_reg)
ax.plot(Xr.ravel(), Yr.ravel(), "ko", ms=8, label="Grid points (vertices)")
ax.axhline(0, color="C3", ls="--", lw=2, label="Boundary")
ax.axhline(L, color="C3", ls="--", lw=2)
ax.axvline(0, color="C3", ls="--", lw=2)
ax.axvline(L, color="C3", ls="--", lw=2)
ax.set_title("Regular (vertex-centred)", fontsize=13)
ax.set_xlim(-0.15, L + 0.15)
ax.set_ylim(-0.15, L + 0.15)
ax.set_aspect("equal")
ax.legend(loc="upper right", fontsize=9)

# Staggered grid
ax = axes[1]
Xs, Ys = np.meshgrid(x_stag, y_stag)
ax.plot(Xs.ravel(), Ys.ravel(), "s", color="C0", ms=8, label="Grid points (cell centres)")
ax.axhline(0, color="C3", ls="--", lw=2, label="Boundary (off-grid)")
ax.axhline(L, color="C3", ls="--", lw=2)
ax.axvline(0, color="C3", ls="--", lw=2)
ax.axvline(L, color="C3", ls="--", lw=2)
ax.set_title("Staggered (cell-centred)", fontsize=13)
ax.set_xlim(-0.15, L + 0.15)
ax.set_ylim(-0.15, L + 0.15)
ax.set_aspect("equal")
ax.legend(loc="upper right", fontsize=9)

plt.suptitle("Grid Geometry: Regular vs Staggered", fontsize=14, y=1.02)
plt.tight_layout()
fig.savefig(IMG_DIR / "grid_geometry.png", dpi=150, bbox_inches="tight")
plt.show()

Grid Geometry

2. Eigenvalue Spectra¶

Each transform type produces a different set of 1-D Laplacian eigenvalues. The formulas are:

Transform	Formula	Null mode?
DST-I	$\lambda_k = -\frac{4}{dx^2}\sin^2\!\left(\frac{\pi(k+1)}{2(N+1)}\right)$	No
DST-II	$\lambda_k = -\frac{4}{dx^2}\sin^2\!\left(\frac{\pi(k+1)}{2N}\right)$	No
DCT-I	$\lambda_k = -\frac{4}{dx^2}\sin^2\!\left(\frac{\pi k}{2(N-1)}\right)$	Yes ($k=0$)
DCT-II	$\lambda_k = -\frac{4}{dx^2}\sin^2\!\left(\frac{\pi k}{2N}\right)$	Yes ($k=0$)
FFT	$\lambda_k = -\frac{4}{dx^2}\sin^2\!\left(\frac{\pi k}{N}\right)$	Yes ($k=0$)

The key difference is the denominator inside the sine: $2(N+1)$, $2N$, or $2(N-1)$. This shifts the eigenvalue curves relative to each other.

In [ ]:

Copied!





N_eig = 32
dx_eig = 1.0
k = np.arange(N_eig)

eig_dst1 = np.array(dst1_eigenvalues(N_eig, dx_eig))
eig_dst2 = np.array(dst2_eigenvalues(N_eig, dx_eig))
eig_dct1 = np.array(dct1_eigenvalues(N_eig, dx_eig))
eig_dct2 = np.array(dct2_eigenvalues(N_eig, dx_eig))
eig_fft = np.array(fft_eigenvalues(N_eig, dx_eig))

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(k, eig_dst1, "o-", ms=4, label="DST-I (Dirichlet, regular)")
ax.plot(k, eig_dst2, "s-", ms=4, label="DST-II (Dirichlet, staggered)")
ax.plot(k, eig_dct1, "^-", ms=4, label="DCT-I (Neumann, regular)")
ax.plot(k, eig_dct2, "v-", ms=4, label="DCT-II (Neumann, staggered)")
ax.plot(k, eig_fft, "D-", ms=4, label="FFT (Periodic)")
ax.set_xlabel("Mode index $k$", fontsize=12)
ax.set_ylabel(r"Eigenvalue $\lambda_k$", fontsize=12)
ax.set_title(f"1-D Laplacian Eigenvalue Spectra ($N={N_eig}$, $dx={dx_eig}$)", fontsize=13)
ax.legend(fontsize=10)
ax.axhline(0, color="k", ls=":", lw=0.5)
ax.grid(True, alpha=0.3)
plt.tight_layout()
fig.savefig(IMG_DIR / "eigenvalue_spectra.png", dpi=150, bbox_inches="tight")
plt.show()
N_eig = 32
dx_eig = 1.0
k = np.arange(N_eig)

eig_dst1 = np.array(dst1_eigenvalues(N_eig, dx_eig))
eig_dst2 = np.array(dst2_eigenvalues(N_eig, dx_eig))
eig_dct1 = np.array(dct1_eigenvalues(N_eig, dx_eig))
eig_dct2 = np.array(dct2_eigenvalues(N_eig, dx_eig))
eig_fft = np.array(fft_eigenvalues(N_eig, dx_eig))

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(k, eig_dst1, "o-", ms=4, label="DST-I (Dirichlet, regular)")
ax.plot(k, eig_dst2, "s-", ms=4, label="DST-II (Dirichlet, staggered)")
ax.plot(k, eig_dct1, "^-", ms=4, label="DCT-I (Neumann, regular)")
ax.plot(k, eig_dct2, "v-", ms=4, label="DCT-II (Neumann, staggered)")
ax.plot(k, eig_fft, "D-", ms=4, label="FFT (Periodic)")
ax.set_xlabel("Mode index $k$", fontsize=12)
ax.set_ylabel(r"Eigenvalue $\lambda_k$", fontsize=12)
ax.set_title(f"1-D Laplacian Eigenvalue Spectra ($N={N_eig}$, $dx={dx_eig}$)", fontsize=13)
ax.legend(fontsize=10)
ax.axhline(0, color="k", ls=":", lw=0.5)
ax.grid(True, alpha=0.3)
plt.tight_layout()
fig.savefig(IMG_DIR / "eigenvalue_spectra.png", dpi=150, bbox_inches="tight")
plt.show()

Eigenvalue Spectra

Observations:

DST-I and DST-II (Dirichlet) are all strictly negative — no null mode.
DCT-I, DCT-II, and FFT have $\lambda_0 = 0$ (constant null mode). For Poisson solves, this mode is projected out (zero-mean gauge).
DST-II eigenvalues are slightly more negative than DST-I at each $k$ because $2N < 2(N+1)$ makes the sine argument larger.
FFT eigenvalues are symmetric about $k = N/2$ (aliasing).

2b. Discrete vs Continuous Eigenvalues¶

The spectral solvers use discrete finite-difference eigenvalues, not the continuous $-k^2$. For low wavenumbers these agree, but they diverge near the Nyquist frequency $k_{\max} = N/2$:

$$ \lambda_k^{\text{discrete}} = -\frac{4}{dx^2}\sin^2\!\left(\frac{\pi k}{2N}\right) \approx -\frac{\pi^2 k^2}{N^2 dx^2} = -k_{\text{phys}}^2 \quad \text{for } k \ll N $$

The discrete eigenvalues saturate at $-4/dx^2$ while the continuous ones grow without bound. This is why spectral solvers are exact inverses of the discrete Laplacian (not approximations to the continuous one).

In [ ]:

Copied!





from spectraldiffx import dst1_eigenvalues, dct2_eigenvalues, fft_eigenvalues

N_dc = 32
dx_dc = 1.0
k_dc = np.arange(N_dc)

eig_dst = np.array(dst1_eigenvalues(N_dc, dx_dc))
eig_dct = np.array(dct2_eigenvalues(N_dc, dx_dc))
eig_fft = np.array(fft_eigenvalues(N_dc, dx_dc))

# Continuous eigenvalues for comparison
cont_dst = -((np.pi * (k_dc + 1) / (N_dc + 1)) ** 2) / dx_dc**2
cont_dct = -((np.pi * k_dc / N_dc) ** 2) / dx_dc**2
cont_fft = -((2 * np.pi * k_dc / (N_dc * dx_dc)) ** 2)

fig, axes = plt.subplots(1, 3, figsize=(14, 4))

for ax, eig, cont, label in [
    (axes[0], eig_dst, cont_dst, "DST-I (Dirichlet)"),
    (axes[1], eig_dct, cont_dct, "DCT-II (Neumann)"),
    (axes[2], eig_fft, cont_fft, "FFT (Periodic)"),
]:
    ax.plot(k_dc, eig, "o-", ms=4, label="Discrete FD", color="C0")
    ax.plot(k_dc, cont, "s--", ms=3, label=r"Continuous $-k^2$", color="C3", alpha=0.7)
    ax.set_xlabel("Mode index $k$")
    ax.set_ylabel(r"$\lambda_k$")
    ax.set_title(label)
    ax.legend(fontsize=9)
    ax.grid(True, alpha=0.3)

plt.suptitle(f"Discrete vs Continuous Eigenvalues ($N = {N_dc}$)", fontsize=13, y=1.02)
plt.tight_layout()
fig.savefig(IMG_DIR / "discrete_vs_continuous_eigenvalues.png", dpi=150, bbox_inches="tight")
plt.show()
from spectraldiffx import dst1_eigenvalues, dct2_eigenvalues, fft_eigenvalues

N_dc = 32
dx_dc = 1.0
k_dc = np.arange(N_dc)

eig_dst = np.array(dst1_eigenvalues(N_dc, dx_dc))
eig_dct = np.array(dct2_eigenvalues(N_dc, dx_dc))
eig_fft = np.array(fft_eigenvalues(N_dc, dx_dc))

# Continuous eigenvalues for comparison
cont_dst = -((np.pi * (k_dc + 1) / (N_dc + 1)) ** 2) / dx_dc**2
cont_dct = -((np.pi * k_dc / N_dc) ** 2) / dx_dc**2
cont_fft = -((2 * np.pi * k_dc / (N_dc * dx_dc)) ** 2)

fig, axes = plt.subplots(1, 3, figsize=(14, 4))

for ax, eig, cont, label in [
    (axes[0], eig_dst, cont_dst, "DST-I (Dirichlet)"),
    (axes[1], eig_dct, cont_dct, "DCT-II (Neumann)"),
    (axes[2], eig_fft, cont_fft, "FFT (Periodic)"),
]:
    ax.plot(k_dc, eig, "o-", ms=4, label="Discrete FD", color="C0")
    ax.plot(k_dc, cont, "s--", ms=3, label=r"Continuous $-k^2$", color="C3", alpha=0.7)
    ax.set_xlabel("Mode index $k$")
    ax.set_ylabel(r"$\lambda_k$")
    ax.set_title(label)
    ax.legend(fontsize=9)
    ax.grid(True, alpha=0.3)

plt.suptitle(f"Discrete vs Continuous Eigenvalues ($N = {N_dc}$)", fontsize=13, y=1.02)
plt.tight_layout()
fig.savefig(IMG_DIR / "discrete_vs_continuous_eigenvalues.png", dpi=150, bbox_inches="tight")
plt.show()

Discrete vs continuous eigenvalues

At low $k$, the discrete and continuous eigenvalues overlap. Near the Nyquist frequency, the discrete eigenvalues saturate at $-4/dx^2$ while the continuous parabola keeps growing. This saturation is why discrete spectral solvers cannot resolve arbitrarily high wavenumbers.

3. Manufactured Solution: Convergence Test¶

We choose manufactured solutions that satisfy the BCs of each solver:

Dirichlet (DST-I, DST-II): $\psi = \sin(2\pi x)\sin(2\pi y)$ — vanishes at $x,y \in \{0,1\}$.
Neumann (DCT-I, DCT-II): $\psi = \cos(2\pi x)\cos(2\pi y)$ — has $\partial\psi/\partial n = 0$ at $x,y \in \{0,1\}$.
Periodic (FFT): $\psi = \sin(2\pi x)\sin(2\pi y)$ — periodic on $[0,1)$.

In all cases: $f = \nabla^2\psi = -8\pi^2\,\psi$

Important: These solvers invert the discrete finite-difference Laplacian ($[1,-2,1]/dx^2$ stencil), not the continuous one. By computing $f$ from the continuous $\nabla^2$, we introduce an $O(dx^2)$ discretisation error. All five solvers therefore converge at second order.

We solve on the unit square $[0,1]^2$ at increasing resolutions and measure the $L^\infty$ error for each solver.

In [ ]:

Copied!





resolutions = [16, 32, 64, 128, 256]

errors = {
    "DST-I (Dirichlet, regular)": [],
    "DST-II (Dirichlet, staggered)": [],
    "DCT-I (Neumann, regular)": [],
    "DCT-II (Neumann, staggered)": [],
    "FFT (Periodic)": [],
}

for N in resolutions:
    # --- Regular grids (DST-I, DCT-I): points at vertices ---
    # DST-I: N interior points, boundary at 0 and L (excluded)
    dx_r = L / (N + 1)
    x_r = jnp.linspace(dx_r, L - dx_r, N)  # interior only
    Xr, Yr = jnp.meshgrid(x_r, x_r)
    psi_exact_r = jnp.sin(2 * jnp.pi * Xr) * jnp.sin(2 * jnp.pi * Yr)
    rhs_r = -8.0 * jnp.pi**2 * psi_exact_r

    psi_dst1 = solve_poisson_dst(rhs_r, dx_r, dx_r)
    errors["DST-I (Dirichlet, regular)"].append(
        float(jnp.max(jnp.abs(psi_dst1 - psi_exact_r)))
    )

    # DCT-I: N points including both boundaries — Neumann-compatible solution
    dx_c1 = L / (N - 1) if N > 1 else L
    x_c1 = jnp.linspace(0, L, N)
    Xc1, Yc1 = jnp.meshgrid(x_c1, x_c1)
    psi_exact_c1 = jnp.cos(2 * jnp.pi * Xc1) * jnp.cos(2 * jnp.pi * Yc1)
    rhs_c1 = -8.0 * jnp.pi**2 * psi_exact_c1

    psi_dct1 = solve_poisson_dct1(rhs_c1, dx_c1, dx_c1)
    errors["DCT-I (Neumann, regular)"].append(
        float(jnp.max(jnp.abs(psi_dct1 - psi_exact_c1)))
    )

    # --- Staggered grids: points at cell centres ---
    dx_s = L / N
    x_s = jnp.linspace(dx_s / 2, L - dx_s / 2, N)
    Xs, Ys = jnp.meshgrid(x_s, x_s)

    # DST-II (Dirichlet): sin solution vanishes at boundary
    psi_exact_dst2 = jnp.sin(2 * jnp.pi * Xs) * jnp.sin(2 * jnp.pi * Ys)
    rhs_dst2 = -8.0 * jnp.pi**2 * psi_exact_dst2

    psi_dst2 = solve_poisson_dst2(rhs_dst2, dx_s, dx_s)
    errors["DST-II (Dirichlet, staggered)"].append(
        float(jnp.max(jnp.abs(psi_dst2 - psi_exact_dst2)))
    )

    # DCT-II (Neumann): cos solution has zero normal derivative at boundary
    psi_exact_dct2 = jnp.cos(2 * jnp.pi * Xs) * jnp.cos(2 * jnp.pi * Ys)
    rhs_dct2 = -8.0 * jnp.pi**2 * psi_exact_dct2

    psi_dct2 = solve_poisson_dct(rhs_dct2, dx_s, dx_s)
    errors["DCT-II (Neumann, staggered)"].append(
        float(jnp.max(jnp.abs(psi_dct2 - psi_exact_dct2)))
    )

    # --- Periodic (FFT) ---
    # Periodic: sin(2*pi*x) is exactly periodic on [0,1)
    dx_p = L / N
    x_p = jnp.linspace(0, L - dx_p, N)
    Xp, Yp = jnp.meshgrid(x_p, x_p)
    psi_exact_p = jnp.sin(2 * jnp.pi * Xp) * jnp.sin(2 * jnp.pi * Yp)
    rhs_p = -8.0 * jnp.pi**2 * psi_exact_p

    psi_fft = solve_poisson_fft(rhs_p, dx_p, dx_p)
    errors["FFT (Periodic)"].append(
        float(jnp.max(jnp.abs(psi_fft - psi_exact_p)))
    )
resolutions = [16, 32, 64, 128, 256]

errors = {
    "DST-I (Dirichlet, regular)": [],
    "DST-II (Dirichlet, staggered)": [],
    "DCT-I (Neumann, regular)": [],
    "DCT-II (Neumann, staggered)": [],
    "FFT (Periodic)": [],
}

for N in resolutions:
    # --- Regular grids (DST-I, DCT-I): points at vertices ---
    # DST-I: N interior points, boundary at 0 and L (excluded)
    dx_r = L / (N + 1)
    x_r = jnp.linspace(dx_r, L - dx_r, N)  # interior only
    Xr, Yr = jnp.meshgrid(x_r, x_r)
    psi_exact_r = jnp.sin(2 * jnp.pi * Xr) * jnp.sin(2 * jnp.pi * Yr)
    rhs_r = -8.0 * jnp.pi**2 * psi_exact_r

    psi_dst1 = solve_poisson_dst(rhs_r, dx_r, dx_r)
    errors["DST-I (Dirichlet, regular)"].append(
        float(jnp.max(jnp.abs(psi_dst1 - psi_exact_r)))
    )

    # DCT-I: N points including both boundaries — Neumann-compatible solution
    dx_c1 = L / (N - 1) if N > 1 else L
    x_c1 = jnp.linspace(0, L, N)
    Xc1, Yc1 = jnp.meshgrid(x_c1, x_c1)
    psi_exact_c1 = jnp.cos(2 * jnp.pi * Xc1) * jnp.cos(2 * jnp.pi * Yc1)
    rhs_c1 = -8.0 * jnp.pi**2 * psi_exact_c1

    psi_dct1 = solve_poisson_dct1(rhs_c1, dx_c1, dx_c1)
    errors["DCT-I (Neumann, regular)"].append(
        float(jnp.max(jnp.abs(psi_dct1 - psi_exact_c1)))
    )

    # --- Staggered grids: points at cell centres ---
    dx_s = L / N
    x_s = jnp.linspace(dx_s / 2, L - dx_s / 2, N)
    Xs, Ys = jnp.meshgrid(x_s, x_s)

    # DST-II (Dirichlet): sin solution vanishes at boundary
    psi_exact_dst2 = jnp.sin(2 * jnp.pi * Xs) * jnp.sin(2 * jnp.pi * Ys)
    rhs_dst2 = -8.0 * jnp.pi**2 * psi_exact_dst2

    psi_dst2 = solve_poisson_dst2(rhs_dst2, dx_s, dx_s)
    errors["DST-II (Dirichlet, staggered)"].append(
        float(jnp.max(jnp.abs(psi_dst2 - psi_exact_dst2)))
    )

    # DCT-II (Neumann): cos solution has zero normal derivative at boundary
    psi_exact_dct2 = jnp.cos(2 * jnp.pi * Xs) * jnp.cos(2 * jnp.pi * Ys)
    rhs_dct2 = -8.0 * jnp.pi**2 * psi_exact_dct2

    psi_dct2 = solve_poisson_dct(rhs_dct2, dx_s, dx_s)
    errors["DCT-II (Neumann, staggered)"].append(
        float(jnp.max(jnp.abs(psi_dct2 - psi_exact_dct2)))
    )

    # --- Periodic (FFT) ---
    # Periodic: sin(2*pi*x) is exactly periodic on [0,1)
    dx_p = L / N
    x_p = jnp.linspace(0, L - dx_p, N)
    Xp, Yp = jnp.meshgrid(x_p, x_p)
    psi_exact_p = jnp.sin(2 * jnp.pi * Xp) * jnp.sin(2 * jnp.pi * Yp)
    rhs_p = -8.0 * jnp.pi**2 * psi_exact_p

    psi_fft = solve_poisson_fft(rhs_p, dx_p, dx_p)
    errors["FFT (Periodic)"].append(
        float(jnp.max(jnp.abs(psi_fft - psi_exact_p)))
    )

In [ ]:

Copied!





fig, ax = plt.subplots(figsize=(9, 6))
markers = ["o", "s", "^", "v", "D"]
for (name, errs), marker in zip(errors.items(), markers):
    ax.loglog(resolutions, errs, f"{marker}-", ms=7, label=name)

# Reference slope: O(dx^2)
N_ref = np.array(resolutions, dtype=float)
ref = 2.0 * (N_ref[0] / N_ref) ** 2 * max(errors["DST-I (Dirichlet, regular)"][0], 1e-6)
ax.loglog(N_ref, ref, "k--", lw=1, alpha=0.5, label=r"$O(N^{-2})$ reference")

ax.set_xlabel("Grid size $N$", fontsize=12)
ax.set_ylabel(r"$L^\infty$ error", fontsize=12)
ax.set_title("Convergence: Manufactured Solution on Unit Square", fontsize=13)
ax.legend(fontsize=9)
ax.grid(True, which="both", alpha=0.3)
plt.tight_layout()
fig.savefig(IMG_DIR / "convergence.png", dpi=150, bbox_inches="tight")
plt.show()
fig, ax = plt.subplots(figsize=(9, 6))
markers = ["o", "s", "^", "v", "D"]
for (name, errs), marker in zip(errors.items(), markers):
    ax.loglog(resolutions, errs, f"{marker}-", ms=7, label=name)

# Reference slope: O(dx^2)
N_ref = np.array(resolutions, dtype=float)
ref = 2.0 * (N_ref[0] / N_ref) ** 2 * max(errors["DST-I (Dirichlet, regular)"][0], 1e-6)
ax.loglog(N_ref, ref, "k--", lw=1, alpha=0.5, label=r"$O(N^{-2})$ reference")

ax.set_xlabel("Grid size $N$", fontsize=12)
ax.set_ylabel(r"$L^\infty$ error", fontsize=12)
ax.set_title("Convergence: Manufactured Solution on Unit Square", fontsize=13)
ax.legend(fontsize=9)
ax.grid(True, which="both", alpha=0.3)
plt.tight_layout()
fig.savefig(IMG_DIR / "convergence.png", dpi=150, bbox_inches="tight")
plt.show()

Convergence

Observations:

All five solvers converge at second order ($O(dx^2)$). This is because each solver inverts the discrete second-order finite-difference Laplacian ($[1,-2,1]/dx^2$ stencil), while the RHS was computed from the continuous $\nabla^2$. The $O(dx^2)$ gap between the discrete and continuous operators dominates the error for all transform types.
If instead we computed the RHS by applying the discrete stencil to the exact solution, the solve would be exact (to machine precision) for all five solvers, because the transform diagonalises that stencil exactly.
Regular and staggered variants converge at the same rate but with slightly different constants.

4. Solution Visualisation¶

We solve $\nabla^2 \psi = f$ for a Gaussian bump RHS and compare the five solvers side by side. The different boundary conditions produce qualitatively different solutions near the domain edges.

Solvability condition: For Neumann and periodic Poisson problems, a solution exists only if $\langle f \rangle = 0$ (the RHS has zero mean). The Gaussian bump has positive mean, so we subtract it for the DCT and FFT cases. The library handles the null mode by zeroing $\hat\psi_0$, but we make the problem well-posed explicitly.

In [ ]:

Copied!





N = 64
dx_s = L / N
x_s = jnp.linspace(dx_s / 2, L - dx_s / 2, N)
Xs, Ys = jnp.meshgrid(x_s, x_s)

# Gaussian bump centred at (0.5, 0.5)
sigma = 0.12
rhs_gauss = jnp.exp(-((Xs - 0.5) ** 2 + (Ys - 0.5) ** 2) / (2 * sigma**2))
# Zero-mean version for Neumann / periodic (solvability condition)
rhs_gauss_zm = rhs_gauss - jnp.mean(rhs_gauss)

# Solve with staggered solvers (cell-centred grids)
psi_dst2_g = solve_poisson_dst2(rhs_gauss, dx_s, dx_s)
psi_dct2_g = solve_poisson_dct(rhs_gauss_zm, dx_s, dx_s)
psi_fft_g = solve_poisson_fft(rhs_gauss_zm, dx_s, dx_s)

# Regular grids need their own coordinate arrays
dx_r_dir = L / (N + 1)
x_r_dir = jnp.linspace(dx_r_dir, L - dx_r_dir, N)
Xrd, Yrd = jnp.meshgrid(x_r_dir, x_r_dir)
rhs_gauss_rd = jnp.exp(-((Xrd - 0.5) ** 2 + (Yrd - 0.5) ** 2) / (2 * sigma**2))
psi_dst1_g = solve_poisson_dst(rhs_gauss_rd, dx_r_dir, dx_r_dir)

dx_r_neu = L / (N - 1)
x_r_neu = jnp.linspace(0, L, N)
Xrn, Yrn = jnp.meshgrid(x_r_neu, x_r_neu)
rhs_gauss_rn = jnp.exp(-((Xrn - 0.5) ** 2 + (Yrn - 0.5) ** 2) / (2 * sigma**2))
rhs_gauss_rn_zm = rhs_gauss_rn - jnp.mean(rhs_gauss_rn)
psi_dct1_g = solve_poisson_dct1(rhs_gauss_rn_zm, dx_r_neu, dx_r_neu)
N = 64
dx_s = L / N
x_s = jnp.linspace(dx_s / 2, L - dx_s / 2, N)
Xs, Ys = jnp.meshgrid(x_s, x_s)

# Gaussian bump centred at (0.5, 0.5)
sigma = 0.12
rhs_gauss = jnp.exp(-((Xs - 0.5) ** 2 + (Ys - 0.5) ** 2) / (2 * sigma**2))
# Zero-mean version for Neumann / periodic (solvability condition)
rhs_gauss_zm = rhs_gauss - jnp.mean(rhs_gauss)

# Solve with staggered solvers (cell-centred grids)
psi_dst2_g = solve_poisson_dst2(rhs_gauss, dx_s, dx_s)
psi_dct2_g = solve_poisson_dct(rhs_gauss_zm, dx_s, dx_s)
psi_fft_g = solve_poisson_fft(rhs_gauss_zm, dx_s, dx_s)

# Regular grids need their own coordinate arrays
dx_r_dir = L / (N + 1)
x_r_dir = jnp.linspace(dx_r_dir, L - dx_r_dir, N)
Xrd, Yrd = jnp.meshgrid(x_r_dir, x_r_dir)
rhs_gauss_rd = jnp.exp(-((Xrd - 0.5) ** 2 + (Yrd - 0.5) ** 2) / (2 * sigma**2))
psi_dst1_g = solve_poisson_dst(rhs_gauss_rd, dx_r_dir, dx_r_dir)

dx_r_neu = L / (N - 1)
x_r_neu = jnp.linspace(0, L, N)
Xrn, Yrn = jnp.meshgrid(x_r_neu, x_r_neu)
rhs_gauss_rn = jnp.exp(-((Xrn - 0.5) ** 2 + (Yrn - 0.5) ** 2) / (2 * sigma**2))
rhs_gauss_rn_zm = rhs_gauss_rn - jnp.mean(rhs_gauss_rn)
psi_dct1_g = solve_poisson_dct1(rhs_gauss_rn_zm, dx_r_neu, dx_r_neu)

In [ ]:

Copied!





fig, axes = plt.subplots(2, 3, figsize=(15, 9))

data = [
    (axes[0, 0], psi_dst1_g, Xrd, Yrd, "DST-I\n(Dirichlet, regular)"),
    (axes[0, 1], psi_dst2_g, Xs, Ys, "DST-II\n(Dirichlet, staggered)"),
    (axes[0, 2], psi_dct1_g, Xrn, Yrn, "DCT-I\n(Neumann, regular)"),
    (axes[1, 0], psi_dct2_g, Xs, Ys, "DCT-II\n(Neumann, staggered)"),
    (axes[1, 1], psi_fft_g, Xs, Ys, "FFT\n(Periodic)"),
]

# Common colour scale
vmin = min(float(jnp.min(d[1])) for d in data)
vmax = max(float(jnp.max(d[1])) for d in data)

for ax, psi, X, Y, title in data:
    im = ax.pcolormesh(
        np.array(X), np.array(Y), np.array(psi),
        cmap="RdBu_r", vmin=vmin, vmax=vmax, shading="auto",
    )
    ax.set_title(title, fontsize=11)
    ax.set_aspect("equal")
    plt.colorbar(im, ax=ax, shrink=0.8)

axes[1, 2].set_visible(False)

plt.suptitle(
    r"Poisson solutions $\nabla^2\psi = f$ for a Gaussian bump RHS",
    fontsize=14, y=1.01,
)
plt.tight_layout()
fig.savefig(IMG_DIR / "solutions.png", dpi=150, bbox_inches="tight")
plt.show()
fig, axes = plt.subplots(2, 3, figsize=(15, 9))

data = [
    (axes[0, 0], psi_dst1_g, Xrd, Yrd, "DST-I\n(Dirichlet, regular)"),
    (axes[0, 1], psi_dst2_g, Xs, Ys, "DST-II\n(Dirichlet, staggered)"),
    (axes[0, 2], psi_dct1_g, Xrn, Yrn, "DCT-I\n(Neumann, regular)"),
    (axes[1, 0], psi_dct2_g, Xs, Ys, "DCT-II\n(Neumann, staggered)"),
    (axes[1, 1], psi_fft_g, Xs, Ys, "FFT\n(Periodic)"),
]

# Common colour scale
vmin = min(float(jnp.min(d[1])) for d in data)
vmax = max(float(jnp.max(d[1])) for d in data)

for ax, psi, X, Y, title in data:
    im = ax.pcolormesh(
        np.array(X), np.array(Y), np.array(psi),
        cmap="RdBu_r", vmin=vmin, vmax=vmax, shading="auto",
    )
    ax.set_title(title, fontsize=11)
    ax.set_aspect("equal")
    plt.colorbar(im, ax=ax, shrink=0.8)

axes[1, 2].set_visible(False)

plt.suptitle(
    r"Poisson solutions $\nabla^2\psi = f$ for a Gaussian bump RHS",
    fontsize=14, y=1.01,
)
plt.tight_layout()
fig.savefig(IMG_DIR / "solutions.png", dpi=150, bbox_inches="tight")
plt.show()

Solutions

Observations:

Dirichlet solutions (DST-I, DST-II) are pulled to zero at the boundary, creating a sharp gradient near the edges.
Neumann solutions (DCT-I, DCT-II) have zero normal derivative at the boundary, so the solution "flattens out" at the edges.
Periodic solution wraps around — the response appears on both sides.
Regular vs staggered variants look very similar at this resolution; the difference is in where exactly the boundary condition is enforced.

5. Error Maps: Regular vs Staggered¶

To see the difference between regular and staggered more clearly, we compare the Dirichlet solvers (DST-I vs DST-II) on the same manufactured solution, evaluated at the appropriate grid points for each.

In [ ]:

Copied!





N = 64

# DST-I (regular): interior points
dx1 = L / (N + 1)
x1 = jnp.linspace(dx1, L - dx1, N)
X1, Y1 = jnp.meshgrid(x1, x1)
psi_exact_1 = jnp.sin(2 * jnp.pi * X1) * jnp.sin(2 * jnp.pi * Y1)
rhs1 = -8.0 * jnp.pi**2 * psi_exact_1
err_dst1 = jnp.abs(solve_poisson_dst(rhs1, dx1, dx1) - psi_exact_1)

# DST-II (staggered): cell centres
dx2 = L / N
x2 = jnp.linspace(dx2 / 2, L - dx2 / 2, N)
X2, Y2 = jnp.meshgrid(x2, x2)
psi_exact_2 = jnp.sin(2 * jnp.pi * X2) * jnp.sin(2 * jnp.pi * Y2)
rhs2 = -8.0 * jnp.pi**2 * psi_exact_2
err_dst2 = jnp.abs(solve_poisson_dst2(rhs2, dx2, dx2) - psi_exact_2)

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

im0 = axes[0].pcolormesh(np.array(X1), np.array(Y1), np.array(err_dst1), cmap="hot_r", shading="auto")
axes[0].set_title("DST-I error (regular)", fontsize=11)
axes[0].set_aspect("equal")
plt.colorbar(im0, ax=axes[0], shrink=0.8, format="%.1e")

im1 = axes[1].pcolormesh(np.array(X2), np.array(Y2), np.array(err_dst2), cmap="hot_r", shading="auto")
axes[1].set_title("DST-II error (staggered)", fontsize=11)
axes[1].set_aspect("equal")
plt.colorbar(im1, ax=axes[1], shrink=0.8, format="%.1e")

# Cross-section through centre
j_mid = N // 2
axes[2].semilogy(np.array(x1), np.array(err_dst1[j_mid, :]), "o-", ms=3, label="DST-I (regular)")
axes[2].semilogy(np.array(x2), np.array(err_dst2[j_mid, :]), "s-", ms=3, label="DST-II (staggered)")
axes[2].set_xlabel("$x$")
axes[2].set_ylabel(r"$|\psi - \psi_{\rm exact}|$")
axes[2].set_title("Error cross-section at $y = 0.5$", fontsize=11)
axes[2].legend()
axes[2].grid(True, alpha=0.3)

plt.suptitle(
    r"Point-wise error: $|\psi_{\rm computed} - \psi_{\rm exact}|$",
    fontsize=13, y=1.02,
)
plt.tight_layout()
fig.savefig(IMG_DIR / "error_maps.png", dpi=150, bbox_inches="tight")
plt.show()
N = 64

# DST-I (regular): interior points
dx1 = L / (N + 1)
x1 = jnp.linspace(dx1, L - dx1, N)
X1, Y1 = jnp.meshgrid(x1, x1)
psi_exact_1 = jnp.sin(2 * jnp.pi * X1) * jnp.sin(2 * jnp.pi * Y1)
rhs1 = -8.0 * jnp.pi**2 * psi_exact_1
err_dst1 = jnp.abs(solve_poisson_dst(rhs1, dx1, dx1) - psi_exact_1)

# DST-II (staggered): cell centres
dx2 = L / N
x2 = jnp.linspace(dx2 / 2, L - dx2 / 2, N)
X2, Y2 = jnp.meshgrid(x2, x2)
psi_exact_2 = jnp.sin(2 * jnp.pi * X2) * jnp.sin(2 * jnp.pi * Y2)
rhs2 = -8.0 * jnp.pi**2 * psi_exact_2
err_dst2 = jnp.abs(solve_poisson_dst2(rhs2, dx2, dx2) - psi_exact_2)

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

im0 = axes[0].pcolormesh(np.array(X1), np.array(Y1), np.array(err_dst1), cmap="hot_r", shading="auto")
axes[0].set_title("DST-I error (regular)", fontsize=11)
axes[0].set_aspect("equal")
plt.colorbar(im0, ax=axes[0], shrink=0.8, format="%.1e")

im1 = axes[1].pcolormesh(np.array(X2), np.array(Y2), np.array(err_dst2), cmap="hot_r", shading="auto")
axes[1].set_title("DST-II error (staggered)", fontsize=11)
axes[1].set_aspect("equal")
plt.colorbar(im1, ax=axes[1], shrink=0.8, format="%.1e")

# Cross-section through centre
j_mid = N // 2
axes[2].semilogy(np.array(x1), np.array(err_dst1[j_mid, :]), "o-", ms=3, label="DST-I (regular)")
axes[2].semilogy(np.array(x2), np.array(err_dst2[j_mid, :]), "s-", ms=3, label="DST-II (staggered)")
axes[2].set_xlabel("$x$")
axes[2].set_ylabel(r"$|\psi - \psi_{\rm exact}|$")
axes[2].set_title("Error cross-section at $y = 0.5$", fontsize=11)
axes[2].legend()
axes[2].grid(True, alpha=0.3)

plt.suptitle(
    r"Point-wise error: $|\psi_{\rm computed} - \psi_{\rm exact}|$",
    fontsize=13, y=1.02,
)
plt.tight_layout()
fig.savefig(IMG_DIR / "error_maps.png", dpi=150, bbox_inches="tight")
plt.show()

Error Maps

6. Summary¶

Question	Answer
Which solver for cell-centred Dirichlet?	`solve_poisson_dst2` (DST-II)
Which solver for vertex-centred Dirichlet?	`solve_poisson_dst` (DST-I)
Which solver for cell-centred Neumann?	`solve_poisson_dct` (DCT-II)
Which solver for vertex-centred Neumann?	`solve_poisson_dct1` (DCT-I)
Which solver for periodic?	`solve_poisson_fft`
Convergence rate?	$O(dx^2)$ for all five (discrete FD Laplacian)
When does the choice matter?	Always — using the wrong grid type introduces $O(dx^2)$ boundary error

The choice of solver must match where your unknowns live on the grid. In finite-volume methods (e.g., Arakawa C-grids), pressure and tracers live at cell centres (staggered), so DST-II/DCT-II are the correct choice. Vorticity/streamfunction at corners use DST-I/DCT-I.