Elliptic Solvers on Different Geometries¶

This notebook demonstrates finitevolX's elliptic solver stack by solving the Helmholtz equation

$$ (\nabla^2 - \lambda)\,\psi = f $$

on four progressively harder domain geometries. For each geometry we solve a manufactured problem, visualise the RHS / solution / residual error triplet, and collect timing data for a final accuracy-vs-speed comparison.

Discrete 5-point stencil¶

On a uniform grid with spacing $\Delta x, \Delta y$ the Laplacian is approximated by the standard 5-point stencil:

$$ \nabla^2 \psi \approx \frac{\psi_{i+1,j} - 2\psi_{i,j} + \psi_{i-1,j}}{\Delta x^2} + \frac{\psi_{i,j+1} - 2\psi_{i,j} + \psi_{i,j-1}}{\Delta y^2} $$

so the discrete Helmholtz operator reads $A\psi = \nabla^2_h \psi - \lambda\,\psi$.

Solver taxonomy¶

Solver	Mechanism	Handles masks?	Variable coeff?	Complexity
Spectral DST	Eigenvalue division in frequency space	No	No	$O(N \log N)$
Capacitance matrix	Sherman-Morrison-Woodbury correction	Yes (few bdry pts)	No	$O(N \log N + B^2)$
CG + spectral PC	Krylov iteration, spectral preconditioner	Yes	No	$O(k \cdot N \log N)$
Multigrid	Recursive V-cycle coarsening / smoothing	Yes	Yes	$O(N)$
MG + CG	MG V-cycle as CG preconditioner	Yes	Yes	$O(k \cdot N)$

In [ ]:

from __future__ import annotations

from pathlib import Path
import time

import jax
import jax.numpy as jnp
import matplotlib

matplotlib.use("Agg")
from matplotlib.patches import Patch
import matplotlib.pyplot as plt
import numpy as np

jax.config.update("jax_enable_x64", True)

import finitevolx as fvx

# Internal import: _apply_operator applies the multigrid's finest operator.
# With coeff=None and the default bc="dirichlet" it equals fvx.masked_laplacian;
# the variable-coefficient notch (section 6) needs it, since masked_laplacian
# has no c(x, y).  There is no public equivalent yet.
from finitevolx._src.solvers.multigrid import _apply_operator

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

from __future__ import annotations from pathlib import Path import time import jax import jax.numpy as jnp import matplotlib matplotlib.use("Agg") from matplotlib.patches import Patch import matplotlib.pyplot as plt import numpy as np jax.config.update("jax_enable_x64", True) import finitevolx as fvx # Internal import: _apply_operator applies the multigrid's finest operator. # With coeff=None and the default bc="dirichlet" it equals fvx.masked_laplacian; # the variable-coefficient notch (section 6) needs it, since masked_laplacian # has no c(x, y). There is no public equivalent yet. from finitevolx._src.solvers.multigrid import _apply_operator IMG_DIR = Path(__file__).resolve().parent.parent / "images" / "demo_solvers" IMG_DIR.mkdir(parents=True, exist_ok=True)

1. Shared Setup¶

Grid parameters¶

We use a 64 x 64 grid on the unit square $[0,1]^2$: small enough for interactive exploration. (Multigrid coarsens any shape, so the size need not be a power of two.) For timings at 128² and 256², with errors against a sparse direct solve, see the Solver Comparison.

Helmholtz parameter¶

We choose $\lambda = 4$. A positive $\lambda$ shifts the operator away from the Poisson null-space, making the system strictly negative-definite and well-conditioned. $\lambda = 4$ is large enough to show the damping effect on the solution amplitude but small enough that the system is still dominated by the Laplacian.

RHS field¶

A single-mode sinusoidal $f(x,y) = \sin(\pi x') \sin(\pi y')$ (where $x', y'$ are normalised to vanish at the boundaries) is smooth, symmetric, and has an analytic eigenfunction expansion — ideal for verifying solvers.

In [ ]:

Ny, Nx = 64, 64
Lx, Ly = 1.0, 1.0
dx, dy = Lx / Nx, Ly / Ny
lambda_ = 4.0

j_idx = jnp.arange(Ny)[:, None]
i_idx = jnp.arange(Nx)[None, :]
rhs = jnp.sin(jnp.pi * (j_idx + 1) / (Ny + 1)) * jnp.sin(
    jnp.pi * (i_idx + 1) / (Nx + 1)
)

Y, X = np.mgrid[:Ny, :Nx]

print(f"Grid: {Ny}x{Nx},  dx={dx:.4f},  dy={dy:.4f},  lambda={lambda_}")

Ny, Nx = 64, 64 Lx, Ly = 1.0, 1.0 dx, dy = Lx / Nx, Ly / Ny lambda_ = 4.0 j_idx = jnp.arange(Ny)[:, None] i_idx = jnp.arange(Nx)[None, :] rhs = jnp.sin(jnp.pi * (j_idx + 1) / (Ny + 1)) * jnp.sin( jnp.pi * (i_idx + 1) / (Nx + 1) ) Y, X = np.mgrid[:Ny, :Nx] print(f"Grid: {Ny}x{Nx}, dx={dx:.4f}, dy={dy:.4f}, lambda={lambda_}")

Plotting and timing helpers¶

plot_triplet produces the standard 3-panel figure used throughout: RHS | Solution | Error, where "Error" is the operator residual $f - A\hat\psi$ restricted to wet cells.

time_fn JIT-compiles a zero-argument callable, warms it up, then averages wall-clock time over several calls.

In [ ]:

def _erode_mask(mask: np.ndarray) -> np.ndarray:
    """Erode a binary mask by 1 cell (min-pool with 3x3 kernel).

    This removes the outermost wet layer — i.e. coast-adjacent cells
    where the discrete Laplacian stencil reads land zeros, producing
    spurious residual values that dominate the error colorbar.
    """
    from scipy.ndimage import minimum_filter

    return minimum_filter(mask, size=3)


def plot_triplet(rhs_field, sol_field, err_field, mask, suptitle, solver_info, fname):
    """Save and show a 3-panel RHS | Solution | Error figure."""
    rhs_np = np.where(mask > 0.5, np.asarray(rhs_field), np.nan)
    sol_np = np.where(mask > 0.5, np.asarray(sol_field), np.nan)

    # For the error panel, erode the mask by 1 cell.  At boundary-adjacent
    # cells the Laplacian stencil reads land zeros, creating O(eps) residual
    # artefacts that visually dominate the colorbar.  Eroding hides this
    # boundary layer so the plot shows the true interior error structure.
    interior_mask = _erode_mask(mask)
    err_np = np.where(interior_mask > 0.5, np.asarray(err_field), np.nan)

    fig, axes = plt.subplots(1, 3, figsize=(13, 3.5))

    im0 = axes[0].imshow(rhs_np, origin="lower", cmap="RdBu_r", interpolation="nearest")
    axes[0].set_title("RHS  $f$", fontsize=11)
    fig.colorbar(im0, ax=axes[0], shrink=0.8)

    im1 = axes[1].imshow(sol_np, origin="lower", cmap="RdBu_r", interpolation="nearest")
    axes[1].set_title(f"Solution  ({solver_info})", fontsize=11)
    fig.colorbar(im1, ax=axes[1], shrink=0.8)

    err_abs_max = np.nanmax(np.abs(err_np)) if np.any(np.isfinite(err_np)) else 1.0
    if err_abs_max == 0.0:
        err_abs_max = 1.0
    im2 = axes[2].imshow(
        err_np,
        origin="lower",
        cmap="RdBu_r",
        interpolation="nearest",
        vmin=-err_abs_max,
        vmax=err_abs_max,
    )
    axes[2].set_title("Error  $f - A\\hat{\\psi}$", fontsize=11)
    fig.colorbar(im2, ax=axes[2], shrink=0.8)

    for ax in axes:
        ax.axis("off")

    fig.suptitle(suptitle, fontsize=14, y=1.02)
    plt.tight_layout()
    fig.savefig(IMG_DIR / fname, dpi=150, bbox_inches="tight")
    plt.show()
    print(f"  Saved {fname}")


def time_fn(fn, warmup=2, repeats=5):
    """JIT-compile, warm up, then time *repeats* calls (seconds)."""
    jitted = jax.jit(fn)
    for _ in range(warmup):
        jitted().block_until_ready()
    t0 = time.perf_counter()
    for _ in range(repeats):
        jitted().block_until_ready()
    return (time.perf_counter() - t0) / repeats


def rel_residual(sol, rhs_loc, matvec, interior_mask=None):
    """||rhs - A(sol)|| / ||rhs||, optionally restricted to interior."""
    residual = rhs_loc - matvec(sol)
    if interior_mask is not None:
        residual = residual * interior_mask
        rhs_loc = rhs_loc * interior_mask
    rhs_norm = float(jnp.linalg.norm(rhs_loc))
    if rhs_norm == 0.0:
        return 0.0
    return float(jnp.linalg.norm(residual)) / rhs_norm


# Results accumulator: geometry → solver → {time_ms, rel_residual, label}
results: dict[str, dict] = {}

def _erode_mask(mask: np.ndarray) -> np.ndarray: """Erode a binary mask by 1 cell (min-pool with 3x3 kernel). This removes the outermost wet layer — i.e. coast-adjacent cells where the discrete Laplacian stencil reads land zeros, producing spurious residual values that dominate the error colorbar. """ from scipy.ndimage import minimum_filter return minimum_filter(mask, size=3) def plot_triplet(rhs_field, sol_field, err_field, mask, suptitle, solver_info, fname): """Save and show a 3-panel RHS | Solution | Error figure.""" rhs_np = np.where(mask > 0.5, np.asarray(rhs_field), np.nan) sol_np = np.where(mask > 0.5, np.asarray(sol_field), np.nan) # For the error panel, erode the mask by 1 cell. At boundary-adjacent # cells the Laplacian stencil reads land zeros, creating O(eps) residual # artefacts that visually dominate the colorbar. Eroding hides this # boundary layer so the plot shows the true interior error structure. interior_mask = _erode_mask(mask) err_np = np.where(interior_mask > 0.5, np.asarray(err_field), np.nan) fig, axes = plt.subplots(1, 3, figsize=(13, 3.5)) im0 = axes[0].imshow(rhs_np, origin="lower", cmap="RdBu_r", interpolation="nearest") axes[0].set_title("RHS $f$", fontsize=11) fig.colorbar(im0, ax=axes[0], shrink=0.8) im1 = axes[1].imshow(sol_np, origin="lower", cmap="RdBu_r", interpolation="nearest") axes[1].set_title(f"Solution ({solver_info})", fontsize=11) fig.colorbar(im1, ax=axes[1], shrink=0.8) err_abs_max = np.nanmax(np.abs(err_np)) if np.any(np.isfinite(err_np)) else 1.0 if err_abs_max == 0.0: err_abs_max = 1.0 im2 = axes[2].imshow( err_np, origin="lower", cmap="RdBu_r", interpolation="nearest", vmin=-err_abs_max, vmax=err_abs_max, ) axes[2].set_title("Error $f - A\\hat{\\psi}$", fontsize=11) fig.colorbar(im2, ax=axes[2], shrink=0.8) for ax in axes: ax.axis("off") fig.suptitle(suptitle, fontsize=14, y=1.02) plt.tight_layout() fig.savefig(IMG_DIR / fname, dpi=150, bbox_inches="tight") plt.show() print(f" Saved {fname}") def time_fn(fn, warmup=2, repeats=5): """JIT-compile, warm up, then time *repeats* calls (seconds).""" jitted = jax.jit(fn) for _ in range(warmup): jitted().block_until_ready() t0 = time.perf_counter() for _ in range(repeats): jitted().block_until_ready() return (time.perf_counter() - t0) / repeats def rel_residual(sol, rhs_loc, matvec, interior_mask=None): """||rhs - A(sol)|| / ||rhs||, optionally restricted to interior.""" residual = rhs_loc - matvec(sol) if interior_mask is not None: residual = residual * interior_mask rhs_loc = rhs_loc * interior_mask rhs_norm = float(jnp.linalg.norm(rhs_loc)) if rhs_norm == 0.0: return 0.0 return float(jnp.linalg.norm(residual)) / rhs_norm # Results accumulator: geometry → solver → {time_ms, rel_residual, label} results: dict[str, dict] = {}

2. Domain Geometries Overview¶

We test four domains of increasing difficulty.

 Rectangle        Basin             Circle            Notch
 ┌──────────┐    ┌──────────┐    ┌──────────┐    ┌──────────┐
 │██████████│    │░░░░░░░░░░│    │    ██    │    │░░░░░░░░░░│
 │██████████│    │░████████░│    │  ██████  │    │░████████░│
 │██████████│    │░████████░│    │ ████████ │    │░████████░│
 │██████████│    │░████████░│    │  ██████  │    │░░░░█████░│
 │██████████│    │░████████░│    │    ██    │    │░░░░█████░│
 └──────────┘    │░░░░░░░░░░│    └──────────┘    │░████████░│
                 └──────────┘                    │░░░░░░░░░░│
                                                 └──────────┘
 ██ = wet (mask=1)   ░ = land (mask=0)

• Rectangle — full grid, no masking. Spectral solver applies directly.
• Basin — rectangular ocean with a 4-cell land border on all sides.
• Circle — disk inscribed in the grid (radius = 0.4 N).
• Notch — rectangle with thick walls and a notch cut from the left, combined with a spatially varying coefficient $c(x,y)$.

In [ ]:

mask_rect = np.ones((Ny, Nx))

mask_basin = np.ones((Ny, Nx))
mask_basin[:4, :] = mask_basin[-4:, :] = 0.0
mask_basin[:, :4] = mask_basin[:, -4:] = 0.0

mask_circle = ((X - Nx / 2) ** 2 + (Y - Ny / 2) ** 2 < (0.4 * Ny) ** 2).astype(float)

mask_notch = np.ones((Ny, Nx))
mask_notch[:8, :] = mask_notch[-8:, :] = 0.0
mask_notch[:, :8] = mask_notch[:, -8:] = 0.0
mask_notch[20:44, :20] = 0.0

geometries = {
    "Rectangle": mask_rect,
    "Basin": mask_basin,
    "Circle": mask_circle,
    "Notch\n(variable coeff)": mask_notch,
}

fig, axes = plt.subplots(1, 4, figsize=(16, 3.5))
for ax, (name, mask) in zip(axes, geometries.items(), strict=True):
    ax.imshow(
        mask, origin="lower", cmap="Blues", interpolation="nearest", vmin=0, vmax=1
    )
    wet = int(mask.sum())
    ax.set_title(f"{name}\n({wet} wet cells)", fontsize=11)
    ax.axis("off")
fig.suptitle("Domain Geometries", fontsize=14, y=1.02)
plt.tight_layout()
fig.savefig(IMG_DIR / "geometries.png", dpi=150, bbox_inches="tight")
plt.show()
print("Saved geometries.png")

mask_rect = np.ones((Ny, Nx)) mask_basin = np.ones((Ny, Nx)) mask_basin[:4, :] = mask_basin[-4:, :] = 0.0 mask_basin[:, :4] = mask_basin[:, -4:] = 0.0 mask_circle = ((X - Nx / 2) ** 2 + (Y - Ny / 2) ** 2 < (0.4 * Ny) ** 2).astype(float) mask_notch = np.ones((Ny, Nx)) mask_notch[:8, :] = mask_notch[-8:, :] = 0.0 mask_notch[:, :8] = mask_notch[:, -8:] = 0.0 mask_notch[20:44, :20] = 0.0 geometries = { "Rectangle": mask_rect, "Basin": mask_basin, "Circle": mask_circle, "Notch\n(variable coeff)": mask_notch, } fig, axes = plt.subplots(1, 4, figsize=(16, 3.5)) for ax, (name, mask) in zip(axes, geometries.items(), strict=True): ax.imshow( mask, origin="lower", cmap="Blues", interpolation="nearest", vmin=0, vmax=1 ) wet = int(mask.sum()) ax.set_title(f"{name}\n({wet} wet cells)", fontsize=11) ax.axis("off") fig.suptitle("Domain Geometries", fontsize=14, y=1.02) plt.tight_layout() fig.savefig(IMG_DIR / "geometries.png", dpi=150, bbox_inches="tight") plt.show() print("Saved geometries.png")

3. Rectangle — Spectral DST, CG, Multigrid¶

The simplest geometry: no land mask, pure Dirichlet boundary conditions. Three solvers compete:

Spectral DST (exact)¶

The Discrete Sine Transform diagonalises the 5-point Laplacian with homogeneous Dirichlet BCs. The Helmholtz solve reduces to pointwise division by known eigenvalues in frequency space:

$$ \hat\psi_{k,l} = \frac{\hat{f}_{k,l}} {\mu_k + \mu_l - \lambda}, \qquad \mu_k = \frac{2}{\Delta x^2}\bigl(\cos(k\pi / (N+1)) - 1\bigr) $$

Cost is $O(N \log N)$ (two DSTs) with zero iteration — the answer is exact up to floating-point rounding.

CG with spectral preconditioner¶

Conjugate Gradient treats the Helmholtz operator as a black-box matrix-vector product (Krylov method). The spectral preconditioner approximately inverts each CG step via a DST solve, giving rapid convergence. fvx.masked_laplacian treats everything beyond the array edge as dry (zero), so on the all-ones rectangle mask it is the Dirichlet operator the DST inverts: CG solves the same problem as the spectral solver, and with this single-mode right-hand side (a DST eigenvector) it converges in a couple of iterations, preconditioned or not. The DST preconditioner is exact here; the FFT one would only be an approximation.

Multigrid¶

Geometric multigrid recursively smooths on a hierarchy of coarsened grids (V-cycles). Each V-cycle costs $O(N)$; 8 cycles suffice for good accuracy on a 64x64 grid.

In [ ]:

print("Rectangle:")
mask_rect_jnp = jnp.array(mask_rect)
interior_rect = jnp.zeros((Ny, Nx)).at[1:-1, 1:-1].set(1.0)

# ── Spectral (exact) ──
sol_sp_rect = fvx.solve_helmholtz_dst(rhs, dx, dy, lambda_)
t_sp_rect = time_fn(lambda: fvx.solve_helmholtz_dst(rhs, dx, dy, lambda_))
A_rect_sp = lambda x: fvx.masked_laplacian(x, mask_rect_jnp, dx, dy, lambda_=lambda_)
err_sp_rect = rhs - A_rect_sp(sol_sp_rect)

plot_triplet(
    rhs,
    sol_sp_rect,
    err_sp_rect,
    mask_rect,
    "Rectangle — Spectral DST (exact)",
    "exact, DST",
    "solver_rect_spectral.png",
)

print("Rectangle:") mask_rect_jnp = jnp.array(mask_rect) interior_rect = jnp.zeros((Ny, Nx)).at[1:-1, 1:-1].set(1.0) # ── Spectral (exact) ── sol_sp_rect = fvx.solve_helmholtz_dst(rhs, dx, dy, lambda_) t_sp_rect = time_fn(lambda: fvx.solve_helmholtz_dst(rhs, dx, dy, lambda_)) A_rect_sp = lambda x: fvx.masked_laplacian(x, mask_rect_jnp, dx, dy, lambda_=lambda_) err_sp_rect = rhs - A_rect_sp(sol_sp_rect) plot_triplet( rhs, sol_sp_rect, err_sp_rect, mask_rect, "Rectangle — Spectral DST (exact)", "exact, DST", "solver_rect_spectral.png", )

In [ ]:

# ── CG (no preconditioner) ──
A_rect = lambda x: fvx.masked_laplacian(x, mask_rect_jnp, dx, dy, lambda_=lambda_)
sol_cg0_rect, info_cg0_rect = fvx.solve_cg(
    A_rect, rhs, rtol=1e-10, atol=1e-10
)
t_cg0_rect = time_fn(
    lambda: fvx.solve_cg(A_rect, rhs, rtol=1e-10, atol=1e-10)[0]
)
rr_cg0_rect = rel_residual(sol_cg0_rect, rhs, A_rect, interior_mask=interior_rect)

# ── CG + spectral preconditioner ──
pc_rect = fvx.make_spectral_preconditioner(dx, dy, lambda_=lambda_, bc="dst")
sol_cg_rect, info_cg_rect = fvx.solve_cg(
    A_rect, rhs, preconditioner=pc_rect, rtol=1e-10, atol=1e-10
)
t_cg_rect = time_fn(
    lambda: fvx.solve_cg(A_rect, rhs, preconditioner=pc_rect, rtol=1e-10, atol=1e-10)[0]
)
rr_cg_rect = rel_residual(sol_cg_rect, rhs, A_rect, interior_mask=interior_rect)
err_cg_rect = rhs - A_rect(sol_cg_rect)

plot_triplet(
    rhs,
    sol_cg_rect,
    err_cg_rect,
    mask_rect,
    f"Rectangle — CG + spectral PC ({info_cg_rect.iterations} iters)",
    f"CG, {info_cg_rect.iterations} iters",
    "solver_rect_cg.png",
)

# ── CG (no preconditioner) ── A_rect = lambda x: fvx.masked_laplacian(x, mask_rect_jnp, dx, dy, lambda_=lambda_) sol_cg0_rect, info_cg0_rect = fvx.solve_cg( A_rect, rhs, rtol=1e-10, atol=1e-10 ) t_cg0_rect = time_fn( lambda: fvx.solve_cg(A_rect, rhs, rtol=1e-10, atol=1e-10)[0] ) rr_cg0_rect = rel_residual(sol_cg0_rect, rhs, A_rect, interior_mask=interior_rect) # ── CG + spectral preconditioner ── pc_rect = fvx.make_spectral_preconditioner(dx, dy, lambda_=lambda_, bc="dst") sol_cg_rect, info_cg_rect = fvx.solve_cg( A_rect, rhs, preconditioner=pc_rect, rtol=1e-10, atol=1e-10 ) t_cg_rect = time_fn( lambda: fvx.solve_cg(A_rect, rhs, preconditioner=pc_rect, rtol=1e-10, atol=1e-10)[0] ) rr_cg_rect = rel_residual(sol_cg_rect, rhs, A_rect, interior_mask=interior_rect) err_cg_rect = rhs - A_rect(sol_cg_rect) plot_triplet( rhs, sol_cg_rect, err_cg_rect, mask_rect, f"Rectangle — CG + spectral PC ({info_cg_rect.iterations} iters)", f"CG, {info_cg_rect.iterations} iters", "solver_rect_cg.png", )

In [ ]:

# ── Multigrid (8 V-cycles) ──
mg_rect = fvx.build_multigrid_solver(mask_rect, dx, dy, lambda_=lambda_, n_cycles=8)
sol_mg_rect = mg_rect(rhs)
t_mg_rect = time_fn(lambda: mg_rect(rhs))
A_mg_rect = lambda u: _apply_operator(u, mg_rect.levels[0])
rr_mg_rect = rel_residual(sol_mg_rect, rhs, A_mg_rect, interior_mask=interior_rect)
err_mg_rect = rhs - A_mg_rect(sol_mg_rect)

plot_triplet(
    rhs,
    sol_mg_rect,
    err_mg_rect,
    mask_rect,
    "Rectangle — Multigrid (8 V-cycles)",
    "MG, 8 V-cycles",
    "solver_rect_mg.png",
)

# ── Multigrid (8 V-cycles) ── mg_rect = fvx.build_multigrid_solver(mask_rect, dx, dy, lambda_=lambda_, n_cycles=8) sol_mg_rect = mg_rect(rhs) t_mg_rect = time_fn(lambda: mg_rect(rhs)) A_mg_rect = lambda u: _apply_operator(u, mg_rect.levels[0]) rr_mg_rect = rel_residual(sol_mg_rect, rhs, A_mg_rect, interior_mask=interior_rect) err_mg_rect = rhs - A_mg_rect(sol_mg_rect) plot_triplet( rhs, sol_mg_rect, err_mg_rect, mask_rect, "Rectangle — Multigrid (8 V-cycles)", "MG, 8 V-cycles", "solver_rect_mg.png", )

In [ ]:

results["Rectangle"] = {
    "Spectral": {"time_ms": t_sp_rect * 1000, "rel_residual": 0.0, "label": "exact"},
    "CG (bare)": {
        "time_ms": t_cg0_rect * 1000,
        "rel_residual": rr_cg0_rect,
        "label": f"{info_cg0_rect.iterations} iters",
    },
    "CG": {
        "time_ms": t_cg_rect * 1000,
        "rel_residual": rr_cg_rect,
        "label": f"{info_cg_rect.iterations} iters",
    },
    "Multigrid": {
        "time_ms": t_mg_rect * 1000,
        "rel_residual": rr_mg_rect,
        "label": "8 V-cyc",
    },
}

print(f"  Spectral: {t_sp_rect * 1000:.2f} ms")
print(f"  CG bare:  {t_cg0_rect * 1000:.2f} ms, {info_cg0_rect.iterations} iters")
print(f"  CG + PC:  {t_cg_rect * 1000:.2f} ms, {info_cg_rect.iterations} iters")
print(f"  MG:       {t_mg_rect * 1000:.2f} ms, rel residual = {rr_mg_rect:.2e}")

results["Rectangle"] = { "Spectral": {"time_ms": t_sp_rect * 1000, "rel_residual": 0.0, "label": "exact"}, "CG (bare)": { "time_ms": t_cg0_rect * 1000, "rel_residual": rr_cg0_rect, "label": f"{info_cg0_rect.iterations} iters", }, "CG": { "time_ms": t_cg_rect * 1000, "rel_residual": rr_cg_rect, "label": f"{info_cg_rect.iterations} iters", }, "Multigrid": { "time_ms": t_mg_rect * 1000, "rel_residual": rr_mg_rect, "label": "8 V-cyc", }, } print(f" Spectral: {t_sp_rect * 1000:.2f} ms") print(f" CG bare: {t_cg0_rect * 1000:.2f} ms, {info_cg0_rect.iterations} iters") print(f" CG + PC: {t_cg_rect * 1000:.2f} ms, {info_cg_rect.iterations} iters") print(f" MG: {t_mg_rect * 1000:.2f} ms, rel residual = {rr_mg_rect:.2e}")

4. Basin with Land Border — Capacitance, CG, Multigrid¶

A rectangular ocean basin with a 4-cell land border on all edges. The interior wet region is slightly smaller than the full grid.

Capacitance matrix (Sherman-Morrison-Woodbury)¶

The capacitance method extends the spectral solver to masked domains. It works by:

1. Solving the full-grid (unmasked) spectral problem
2. Precomputing Green's functions at each boundary point $b \in B$
3. Correcting with a dense $B \times B$ capacitance system

The offline cost is $O(B \cdot N \log N)$ for $B$ boundary points plus an $O(B^3)$ factorisation. Each subsequent solve is $O(N \log N + B^2)$. This is fast when $B \ll N$ (near-rectangular masks).

CG + spectral preconditioner¶

Same Krylov approach as for the rectangle, but using fvx.masked_laplacian which zeroes out land cells.

Multigrid¶

Geometric multigrid handles the mask natively by restricting smoothing and residual computation to wet cells.

In [ ]:

print("\nBasin:")
mask_basin_jnp = jnp.array(mask_basin)
rhs_basin = rhs * mask_basin_jnp
A_basin = lambda x: fvx.masked_laplacian(x, mask_basin_jnp, dx, dy, lambda_=lambda_)

# ── Capacitance ──
cap = fvx.build_capacitance_solver(
    mask_basin.astype(bool), dx, dy, lambda_=lambda_, base_bc="dst"
)
sol_cap = cap(rhs_basin)
t_cap = time_fn(lambda: cap(rhs_basin))
err_cap = rhs_basin - A_basin(sol_cap)

plot_triplet(
    rhs_basin,
    sol_cap,
    err_cap,
    mask_basin,
    "Basin — Capacitance Matrix (direct)",
    "capacitance, direct",
    "solver_basin_cap.png",
)
print(f"  Capacitance: {t_cap * 1000:.2f} ms, coast pts = {int(fvx.coast_ring(mask_basin_jnp).sum())}")

print("\nBasin:") mask_basin_jnp = jnp.array(mask_basin) rhs_basin = rhs * mask_basin_jnp A_basin = lambda x: fvx.masked_laplacian(x, mask_basin_jnp, dx, dy, lambda_=lambda_) # ── Capacitance ── cap = fvx.build_capacitance_solver( mask_basin.astype(bool), dx, dy, lambda_=lambda_, base_bc="dst" ) sol_cap = cap(rhs_basin) t_cap = time_fn(lambda: cap(rhs_basin)) err_cap = rhs_basin - A_basin(sol_cap) plot_triplet( rhs_basin, sol_cap, err_cap, mask_basin, "Basin — Capacitance Matrix (direct)", "capacitance, direct", "solver_basin_cap.png", ) print(f" Capacitance: {t_cap * 1000:.2f} ms, coast pts = {int(fvx.coast_ring(mask_basin_jnp).sum())}")

In [ ]:

# ── CG (no preconditioner) ──
sol_cg0_basin, info_cg0_basin = fvx.solve_cg(
    A_basin, rhs_basin, rtol=1e-10, atol=1e-10
)
sol_cg0_basin = sol_cg0_basin * mask_basin_jnp
t_cg0_basin = time_fn(
    lambda: fvx.solve_cg(A_basin, rhs_basin, rtol=1e-10, atol=1e-10)[0] * mask_basin_jnp
)
rr_cg0_basin = rel_residual(
    sol_cg0_basin, rhs_basin, A_basin, interior_mask=mask_basin_jnp
)

# ── CG + spectral preconditioner ──
pc_basin = fvx.make_spectral_preconditioner(dx, dy, lambda_=lambda_, bc="dst")
sol_cg_basin, info_cg_basin = fvx.solve_cg(
    A_basin, rhs_basin, preconditioner=pc_basin, rtol=1e-10, atol=1e-10
)
sol_cg_basin = sol_cg_basin * mask_basin_jnp
t_cg_basin = time_fn(
    lambda: (
        fvx.solve_cg(
            A_basin, rhs_basin, preconditioner=pc_basin, rtol=1e-10, atol=1e-10
        )[0]
        * mask_basin_jnp
    )
)
rr_cg_basin = rel_residual(
    sol_cg_basin, rhs_basin, A_basin, interior_mask=mask_basin_jnp
)
err_cg_basin = (rhs_basin - A_basin(sol_cg_basin)) * mask_basin_jnp

plot_triplet(
    rhs_basin,
    sol_cg_basin,
    err_cg_basin,
    mask_basin,
    f"Basin — CG + spectral PC ({info_cg_basin.iterations} iters)",
    f"CG, {info_cg_basin.iterations} iters",
    "solver_basin_cg.png",
)

# ── CG (no preconditioner) ── sol_cg0_basin, info_cg0_basin = fvx.solve_cg( A_basin, rhs_basin, rtol=1e-10, atol=1e-10 ) sol_cg0_basin = sol_cg0_basin * mask_basin_jnp t_cg0_basin = time_fn( lambda: fvx.solve_cg(A_basin, rhs_basin, rtol=1e-10, atol=1e-10)[0] * mask_basin_jnp ) rr_cg0_basin = rel_residual( sol_cg0_basin, rhs_basin, A_basin, interior_mask=mask_basin_jnp ) # ── CG + spectral preconditioner ── pc_basin = fvx.make_spectral_preconditioner(dx, dy, lambda_=lambda_, bc="dst") sol_cg_basin, info_cg_basin = fvx.solve_cg( A_basin, rhs_basin, preconditioner=pc_basin, rtol=1e-10, atol=1e-10 ) sol_cg_basin = sol_cg_basin * mask_basin_jnp t_cg_basin = time_fn( lambda: ( fvx.solve_cg( A_basin, rhs_basin, preconditioner=pc_basin, rtol=1e-10, atol=1e-10 )[0] * mask_basin_jnp ) ) rr_cg_basin = rel_residual( sol_cg_basin, rhs_basin, A_basin, interior_mask=mask_basin_jnp ) err_cg_basin = (rhs_basin - A_basin(sol_cg_basin)) * mask_basin_jnp plot_triplet( rhs_basin, sol_cg_basin, err_cg_basin, mask_basin, f"Basin — CG + spectral PC ({info_cg_basin.iterations} iters)", f"CG, {info_cg_basin.iterations} iters", "solver_basin_cg.png", )

In [ ]:

# ── Multigrid (8 V-cycles) ──
mg_basin = fvx.build_multigrid_solver(mask_basin, dx, dy, lambda_=lambda_, n_cycles=8)
sol_mg_basin = mg_basin(rhs_basin)
t_mg_basin = time_fn(lambda: mg_basin(rhs_basin))
A_mg_basin = lambda u: _apply_operator(u, mg_basin.levels[0])
rr_mg_basin = rel_residual(
    sol_mg_basin, rhs_basin, A_mg_basin, interior_mask=mask_basin_jnp
)
err_mg_basin = (rhs_basin - A_mg_basin(sol_mg_basin)) * mask_basin_jnp

plot_triplet(
    rhs_basin,
    sol_mg_basin,
    err_mg_basin,
    mask_basin,
    "Basin — Multigrid (8 V-cycles)",
    "MG, 8 V-cycles",
    "solver_basin_mg.png",
)

# ── Multigrid (8 V-cycles) ── mg_basin = fvx.build_multigrid_solver(mask_basin, dx, dy, lambda_=lambda_, n_cycles=8) sol_mg_basin = mg_basin(rhs_basin) t_mg_basin = time_fn(lambda: mg_basin(rhs_basin)) A_mg_basin = lambda u: _apply_operator(u, mg_basin.levels[0]) rr_mg_basin = rel_residual( sol_mg_basin, rhs_basin, A_mg_basin, interior_mask=mask_basin_jnp ) err_mg_basin = (rhs_basin - A_mg_basin(sol_mg_basin)) * mask_basin_jnp plot_triplet( rhs_basin, sol_mg_basin, err_mg_basin, mask_basin, "Basin — Multigrid (8 V-cycles)", "MG, 8 V-cycles", "solver_basin_mg.png", )

In [ ]:

results["Basin"] = {
    "Capacitance": {"time_ms": t_cap * 1000, "rel_residual": 0.0, "label": "direct"},
    "CG (bare)": {
        "time_ms": t_cg0_basin * 1000,
        "rel_residual": rr_cg0_basin,
        "label": f"{info_cg0_basin.iterations} iters",
    },
    "CG": {
        "time_ms": t_cg_basin * 1000,
        "rel_residual": rr_cg_basin,
        "label": f"{info_cg_basin.iterations} iters",
    },
    "Multigrid": {
        "time_ms": t_mg_basin * 1000,
        "rel_residual": rr_mg_basin,
        "label": "8 V-cyc",
    },
}

print(f"  CG bare:  {t_cg0_basin * 1000:.2f} ms, {info_cg0_basin.iterations} iters")
print(f"  CG + PC:  {t_cg_basin * 1000:.2f} ms, {info_cg_basin.iterations} iters")
print(f"  MG:       {t_mg_basin * 1000:.2f} ms, rel residual = {rr_mg_basin:.2e}")

results["Basin"] = { "Capacitance": {"time_ms": t_cap * 1000, "rel_residual": 0.0, "label": "direct"}, "CG (bare)": { "time_ms": t_cg0_basin * 1000, "rel_residual": rr_cg0_basin, "label": f"{info_cg0_basin.iterations} iters", }, "CG": { "time_ms": t_cg_basin * 1000, "rel_residual": rr_cg_basin, "label": f"{info_cg_basin.iterations} iters", }, "Multigrid": { "time_ms": t_mg_basin * 1000, "rel_residual": rr_mg_basin, "label": "8 V-cyc", }, } print(f" CG bare: {t_cg0_basin * 1000:.2f} ms, {info_cg0_basin.iterations} iters") print(f" CG + PC: {t_cg_basin * 1000:.2f} ms, {info_cg_basin.iterations} iters") print(f" MG: {t_mg_basin * 1000:.2f} ms, rel residual = {rr_mg_basin:.2e}")

5. Circular Basin — CG, Multigrid¶

A circular ocean basin (radius = 0.4 N cells) inscribed in the grid. The capacitance method still applies (its setup grows with the number of coast points $B$, which scales with the circumference; the Solver Comparison measures it on a circular basin with an island); here we compare the iterative solvers:

CG with spectral preconditioner (Krylov method)¶

Conjugate Gradient iterates $x_{k+1} = x_k + \alpha_k p_k$ in a Krylov subspace $\mathcal{K}_k(A, r_0)$. The spectral preconditioner applies an approximate inverse via a full-grid DST solve at each step, reducing the effective condition number.

Multigrid standalone¶

The multigrid hierarchy coarsens the circular mask along with the grid. At each level, smoothing and restriction/prolongation respect the mask boundaries, so no CG wrapper is needed.

In [ ]:

print("\nCircle:")
mask_circle_jnp = jnp.array(mask_circle)
rhs_circle = rhs * mask_circle_jnp
A_circle = lambda x: fvx.masked_laplacian(x, mask_circle_jnp, dx, dy, lambda_=lambda_)

# ── CG (no preconditioner) ──
sol_cg0_circle, info_cg0_circle = fvx.solve_cg(
    A_circle, rhs_circle, rtol=1e-10, atol=1e-10
)
sol_cg0_circle = sol_cg0_circle * mask_circle_jnp
t_cg0_circle = time_fn(
    lambda: fvx.solve_cg(A_circle, rhs_circle, rtol=1e-10, atol=1e-10)[0]
    * mask_circle_jnp
)
rr_cg0_circle = rel_residual(
    sol_cg0_circle, rhs_circle, A_circle, interior_mask=mask_circle_jnp
)

# ── CG + spectral preconditioner ──
pc_circle = fvx.make_spectral_preconditioner(dx, dy, lambda_=lambda_, bc="dst")
sol_cg_circle, info_cg_circle = fvx.solve_cg(
    A_circle, rhs_circle, preconditioner=pc_circle, rtol=1e-10, atol=1e-10
)
sol_cg_circle = sol_cg_circle * mask_circle_jnp
t_cg_circle = time_fn(
    lambda: (
        fvx.solve_cg(
            A_circle, rhs_circle, preconditioner=pc_circle, rtol=1e-10, atol=1e-10
        )[0]
        * mask_circle_jnp
    )
)
rr_cg_circle = rel_residual(
    sol_cg_circle, rhs_circle, A_circle, interior_mask=mask_circle_jnp
)
err_cg_circle = (rhs_circle - A_circle(sol_cg_circle)) * mask_circle_jnp

plot_triplet(
    rhs_circle,
    sol_cg_circle,
    err_cg_circle,
    mask_circle,
    f"Circle — CG + spectral PC ({info_cg_circle.iterations} iters)",
    f"CG, {info_cg_circle.iterations} iters",
    "solver_circle_cg.png",
)

print("\nCircle:") mask_circle_jnp = jnp.array(mask_circle) rhs_circle = rhs * mask_circle_jnp A_circle = lambda x: fvx.masked_laplacian(x, mask_circle_jnp, dx, dy, lambda_=lambda_) # ── CG (no preconditioner) ── sol_cg0_circle, info_cg0_circle = fvx.solve_cg( A_circle, rhs_circle, rtol=1e-10, atol=1e-10 ) sol_cg0_circle = sol_cg0_circle * mask_circle_jnp t_cg0_circle = time_fn( lambda: fvx.solve_cg(A_circle, rhs_circle, rtol=1e-10, atol=1e-10)[0] * mask_circle_jnp ) rr_cg0_circle = rel_residual( sol_cg0_circle, rhs_circle, A_circle, interior_mask=mask_circle_jnp ) # ── CG + spectral preconditioner ── pc_circle = fvx.make_spectral_preconditioner(dx, dy, lambda_=lambda_, bc="dst") sol_cg_circle, info_cg_circle = fvx.solve_cg( A_circle, rhs_circle, preconditioner=pc_circle, rtol=1e-10, atol=1e-10 ) sol_cg_circle = sol_cg_circle * mask_circle_jnp t_cg_circle = time_fn( lambda: ( fvx.solve_cg( A_circle, rhs_circle, preconditioner=pc_circle, rtol=1e-10, atol=1e-10 )[0] * mask_circle_jnp ) ) rr_cg_circle = rel_residual( sol_cg_circle, rhs_circle, A_circle, interior_mask=mask_circle_jnp ) err_cg_circle = (rhs_circle - A_circle(sol_cg_circle)) * mask_circle_jnp plot_triplet( rhs_circle, sol_cg_circle, err_cg_circle, mask_circle, f"Circle — CG + spectral PC ({info_cg_circle.iterations} iters)", f"CG, {info_cg_circle.iterations} iters", "solver_circle_cg.png", )

In [ ]:

# ── Multigrid (8 V-cycles) ──
mg_circle = fvx.build_multigrid_solver(mask_circle, dx, dy, lambda_=lambda_, n_cycles=8)
sol_mg_circle = mg_circle(rhs_circle)
t_mg_circle = time_fn(lambda: mg_circle(rhs_circle))
A_mg_circle = lambda u: _apply_operator(u, mg_circle.levels[0])
rr_mg_circle = rel_residual(
    sol_mg_circle, rhs_circle, A_mg_circle, interior_mask=mask_circle_jnp
)
err_mg_circle = (rhs_circle - A_mg_circle(sol_mg_circle)) * mask_circle_jnp

plot_triplet(
    rhs_circle,
    sol_mg_circle,
    err_mg_circle,
    mask_circle,
    "Circle — Multigrid (8 V-cycles)",
    "MG, 8 V-cycles",
    "solver_circle_mg.png",
)

# ── Multigrid (8 V-cycles) ── mg_circle = fvx.build_multigrid_solver(mask_circle, dx, dy, lambda_=lambda_, n_cycles=8) sol_mg_circle = mg_circle(rhs_circle) t_mg_circle = time_fn(lambda: mg_circle(rhs_circle)) A_mg_circle = lambda u: _apply_operator(u, mg_circle.levels[0]) rr_mg_circle = rel_residual( sol_mg_circle, rhs_circle, A_mg_circle, interior_mask=mask_circle_jnp ) err_mg_circle = (rhs_circle - A_mg_circle(sol_mg_circle)) * mask_circle_jnp plot_triplet( rhs_circle, sol_mg_circle, err_mg_circle, mask_circle, "Circle — Multigrid (8 V-cycles)", "MG, 8 V-cycles", "solver_circle_mg.png", )

In [ ]:

results["Circle"] = {
    "CG (bare)": {
        "time_ms": t_cg0_circle * 1000,
        "rel_residual": rr_cg0_circle,
        "label": f"{info_cg0_circle.iterations} iters",
    },
    "CG": {
        "time_ms": t_cg_circle * 1000,
        "rel_residual": rr_cg_circle,
        "label": f"{info_cg_circle.iterations} iters",
    },
    "Multigrid": {
        "time_ms": t_mg_circle * 1000,
        "rel_residual": rr_mg_circle,
        "label": "8 V-cyc",
    },
}

print(f"  CG bare: {t_cg0_circle * 1000:.2f} ms, {info_cg0_circle.iterations} iters")
print(f"  CG + PC: {t_cg_circle * 1000:.2f} ms, {info_cg_circle.iterations} iters")
print(f"  MG:      {t_mg_circle * 1000:.2f} ms, rel residual = {rr_mg_circle:.2e}")

results["Circle"] = { "CG (bare)": { "time_ms": t_cg0_circle * 1000, "rel_residual": rr_cg0_circle, "label": f"{info_cg0_circle.iterations} iters", }, "CG": { "time_ms": t_cg_circle * 1000, "rel_residual": rr_cg_circle, "label": f"{info_cg_circle.iterations} iters", }, "Multigrid": { "time_ms": t_mg_circle * 1000, "rel_residual": rr_mg_circle, "label": "8 V-cyc", }, } print(f" CG bare: {t_cg0_circle * 1000:.2f} ms, {info_cg0_circle.iterations} iters") print(f" CG + PC: {t_cg_circle * 1000:.2f} ms, {info_cg_circle.iterations} iters") print(f" MG: {t_mg_circle * 1000:.2f} ms, rel residual = {rr_mg_circle:.2e}")

6. Notch Domain with Variable Coefficient — Multigrid, MG+CG¶

The hardest case: a spatially varying diffusion coefficient $c(x,y)$ on an irregular masked domain. Only multigrid handles both features natively.

We solve the variable-coefficient Helmholtz equation:

$$ \nabla \cdot \bigl(c(x,y)\,\nabla \psi\bigr) - \lambda\,\psi = f $$

where $c(x,y) = 1 + 0.8\sin(2\pi x / N_x)$ varies sinusoidally in $x$.

The domain is a rectangle with 8-cell walls on all sides and a rectangular notch (rows 20-44, columns 0-20) cut from the left:

  ░░░░░░░░░░░░░░░░░░░░
  ░████████████████████░
  ░████████████████████░
  ░░░░░░███████████████░
  ░░░░░░███████████████░
  ░████████████████████░
  ░████████████████████░
  ░░░░░░░░░░░░░░░░░░░░░

Multigrid standalone (10 V-cycles)¶

The variable coefficient $c(x,y)$ enters the operator through face-averaged diffusivities. Multigrid coarsens both the mask and the coefficient field. We use 10 V-cycles (more than the 8 used for constant-coefficient cases) because the variable coefficient increases the condition number.

MG-preconditioned CG¶

For maximum accuracy, use the multigrid solver as a preconditioner inside CG. Each CG step applies one multigrid V-cycle as an approximate inverse, and CG refines to the requested tolerance.

Important: The CG operator must be the variable-coefficient operator the multigrid was built with (_apply_operator), not fvx.masked_laplacian, which has no $c(x, y)$. (With a constant coefficient the two are the same operator.)

In [ ]:

print("\nNotch (variable coeff):")
coeff = 1.0 + 0.8 * np.sin(2 * np.pi * X / Nx)
mask_notch_jnp = jnp.array(mask_notch)
rhs_notch = rhs * mask_notch_jnp

mg_notch = fvx.build_multigrid_solver(
    mask_notch, dx, dy, lambda_=lambda_, coeff=coeff, n_cycles=10
)
A_notch = lambda x: _apply_operator(x, mg_notch.levels[0])

# ── CG (no preconditioner) ──
sol_cg0_notch, info_cg0_notch = fvx.solve_cg(
    A_notch, rhs_notch, rtol=1e-10, atol=1e-10
)
sol_cg0_notch = sol_cg0_notch * mask_notch_jnp
t_cg0_notch = time_fn(
    lambda: fvx.solve_cg(A_notch, rhs_notch, rtol=1e-10, atol=1e-10)[0]
    * mask_notch_jnp
)
rr_cg0_notch = rel_residual(
    sol_cg0_notch, rhs_notch, A_notch, interior_mask=mask_notch_jnp
)

# ── Multigrid standalone (10 V-cycles) ──
sol_mg_notch = mg_notch(rhs_notch)
t_mg_notch = time_fn(lambda: mg_notch(rhs_notch))
rr_mg_notch = rel_residual(
    sol_mg_notch, rhs_notch, A_notch, interior_mask=mask_notch_jnp
)
err_mg_notch = (rhs_notch - A_notch(sol_mg_notch)) * mask_notch_jnp

plot_triplet(
    rhs_notch,
    sol_mg_notch,
    err_mg_notch,
    mask_notch,
    "Notch — Multigrid standalone (10 V-cycles)",
    "MG, 10 V-cycles",
    "solver_notch_mg.png",
)

print("\nNotch (variable coeff):") coeff = 1.0 + 0.8 * np.sin(2 * np.pi * X / Nx) mask_notch_jnp = jnp.array(mask_notch) rhs_notch = rhs * mask_notch_jnp mg_notch = fvx.build_multigrid_solver( mask_notch, dx, dy, lambda_=lambda_, coeff=coeff, n_cycles=10 ) A_notch = lambda x: _apply_operator(x, mg_notch.levels[0]) # ── CG (no preconditioner) ── sol_cg0_notch, info_cg0_notch = fvx.solve_cg( A_notch, rhs_notch, rtol=1e-10, atol=1e-10 ) sol_cg0_notch = sol_cg0_notch * mask_notch_jnp t_cg0_notch = time_fn( lambda: fvx.solve_cg(A_notch, rhs_notch, rtol=1e-10, atol=1e-10)[0] * mask_notch_jnp ) rr_cg0_notch = rel_residual( sol_cg0_notch, rhs_notch, A_notch, interior_mask=mask_notch_jnp ) # ── Multigrid standalone (10 V-cycles) ── sol_mg_notch = mg_notch(rhs_notch) t_mg_notch = time_fn(lambda: mg_notch(rhs_notch)) rr_mg_notch = rel_residual( sol_mg_notch, rhs_notch, A_notch, interior_mask=mask_notch_jnp ) err_mg_notch = (rhs_notch - A_notch(sol_mg_notch)) * mask_notch_jnp plot_triplet( rhs_notch, sol_mg_notch, err_mg_notch, mask_notch, "Notch — Multigrid standalone (10 V-cycles)", "MG, 10 V-cycles", "solver_notch_mg.png", )

In [ ]:

# ── MG-preconditioned CG ──
pc_mg_notch = fvx.make_multigrid_preconditioner(mg_notch)
sol_mgcg_notch, info_mgcg_notch = fvx.solve_cg(
    A_notch, rhs_notch, preconditioner=pc_mg_notch, rtol=1e-12, atol=1e-12
)
sol_mgcg_notch = sol_mgcg_notch * mask_notch_jnp
t_mgcg_notch = time_fn(
    lambda: (
        fvx.solve_cg(
            A_notch, rhs_notch, preconditioner=pc_mg_notch, rtol=1e-12, atol=1e-12
        )[0]
        * mask_notch_jnp
    )
)
rr_mgcg_notch = rel_residual(
    sol_mgcg_notch, rhs_notch, A_notch, interior_mask=mask_notch_jnp
)
err_mgcg_notch = (rhs_notch - A_notch(sol_mgcg_notch)) * mask_notch_jnp

plot_triplet(
    rhs_notch,
    sol_mgcg_notch,
    err_mgcg_notch,
    mask_notch,
    f"Notch — MG+CG ({info_mgcg_notch.iterations} iters)",
    f"MG+CG, {info_mgcg_notch.iterations} iters",
    "solver_notch_mgcg.png",
)

# ── MG-preconditioned CG ── pc_mg_notch = fvx.make_multigrid_preconditioner(mg_notch) sol_mgcg_notch, info_mgcg_notch = fvx.solve_cg( A_notch, rhs_notch, preconditioner=pc_mg_notch, rtol=1e-12, atol=1e-12 ) sol_mgcg_notch = sol_mgcg_notch * mask_notch_jnp t_mgcg_notch = time_fn( lambda: ( fvx.solve_cg( A_notch, rhs_notch, preconditioner=pc_mg_notch, rtol=1e-12, atol=1e-12 )[0] * mask_notch_jnp ) ) rr_mgcg_notch = rel_residual( sol_mgcg_notch, rhs_notch, A_notch, interior_mask=mask_notch_jnp ) err_mgcg_notch = (rhs_notch - A_notch(sol_mgcg_notch)) * mask_notch_jnp plot_triplet( rhs_notch, sol_mgcg_notch, err_mgcg_notch, mask_notch, f"Notch — MG+CG ({info_mgcg_notch.iterations} iters)", f"MG+CG, {info_mgcg_notch.iterations} iters", "solver_notch_mgcg.png", )

In [ ]:

results["Notch\n(variable coeff)"] = {
    "CG (bare)": {
        "time_ms": t_cg0_notch * 1000,
        "rel_residual": rr_cg0_notch,
        "label": f"{info_cg0_notch.iterations} iters",
    },
    "Multigrid": {
        "time_ms": t_mg_notch * 1000,
        "rel_residual": rr_mg_notch,
        "label": "10 V-cyc",
    },
    "MG+CG": {
        "time_ms": t_mgcg_notch * 1000,
        "rel_residual": rr_mgcg_notch,
        "label": f"{info_mgcg_notch.iterations} iters",
    },
}

diff_mg_mgcg = float(
    jnp.linalg.norm(sol_mgcg_notch - sol_mg_notch) / jnp.linalg.norm(sol_mgcg_notch)
)
print(f"  CG bare: {t_cg0_notch * 1000:.2f} ms, {info_cg0_notch.iterations} iters")
print(f"  MG:      {t_mg_notch * 1000:.2f} ms, rel residual = {rr_mg_notch:.2e}")
print(
    f"  MG+CG:   {t_mgcg_notch * 1000:.2f} ms, {info_mgcg_notch.iterations} iters, "
    f"rel residual = {rr_mgcg_notch:.2e}"
)
print(f"  Relative difference MG vs MG+CG: {diff_mg_mgcg:.2e}")

results["Notch\n(variable coeff)"] = { "CG (bare)": { "time_ms": t_cg0_notch * 1000, "rel_residual": rr_cg0_notch, "label": f"{info_cg0_notch.iterations} iters", }, "Multigrid": { "time_ms": t_mg_notch * 1000, "rel_residual": rr_mg_notch, "label": "10 V-cyc", }, "MG+CG": { "time_ms": t_mgcg_notch * 1000, "rel_residual": rr_mgcg_notch, "label": f"{info_mgcg_notch.iterations} iters", }, } diff_mg_mgcg = float( jnp.linalg.norm(sol_mgcg_notch - sol_mg_notch) / jnp.linalg.norm(sol_mgcg_notch) ) print(f" CG bare: {t_cg0_notch * 1000:.2f} ms, {info_cg0_notch.iterations} iters") print(f" MG: {t_mg_notch * 1000:.2f} ms, rel residual = {rr_mg_notch:.2e}") print( f" MG+CG: {t_mgcg_notch * 1000:.2f} ms, {info_mgcg_notch.iterations} iters, " f"rel residual = {rr_mgcg_notch:.2e}" ) print(f" Relative difference MG vs MG+CG: {diff_mg_mgcg:.2e}")

7. Accuracy & Timing Comparison¶

The bar chart below compares relative residual (left, log scale) and JIT-compiled solve time (right) across all geometry / solver combinations. Timing uses the warmup + repeated-call approach: each solver is JIT-compiled, called twice to warm up, then timed over 5 repeated calls with block_until_ready().

In [ ]:

print("\nSummary chart:")
geom_order = list(results.keys())
geom_shorts = [g.split("\n")[0] for g in geom_order]

colors = {
    "Spectral": "#2196F3",
    "Capacitance": "#4CAF50",
    "CG (bare)": "#BDBDBD",
    "CG": "#FF9800",
    "Multigrid": "#9C27B0",
    "MG+CG": "#E91E63",
}

x_positions: list[float] = []
bar_colors: list[str] = []
times_list: list[float] = []
residuals_list: list[float] = []
solver_names: list[str] = []
group_width = 0.8

for gi, geom in enumerate(geom_order):
    solvers = list(results[geom].keys())
    n_s = len(solvers)
    bar_w = group_width / n_s
    for si, sname in enumerate(solvers):
        x_positions.append(gi + (si - (n_s - 1) / 2) * bar_w)
        bar_colors.append(colors.get(sname, "#888888"))
        times_list.append(results[geom][sname]["time_ms"])
        residuals_list.append(results[geom][sname]["rel_residual"])
        solver_names.append(sname)

fig, (ax_res, ax_time) = plt.subplots(1, 2, figsize=(16, 5))

# ── Timing bars ──
bars = ax_time.bar(
    x_positions,
    times_list,
    width=group_width / 3.5,
    color=bar_colors,
    edgecolor="black",
    linewidth=0.5,
)
for bar, t in zip(bars, times_list, strict=True):
    ax_time.text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + 0.15,
        f"{t:.1f}",
        ha="center",
        va="bottom",
        fontsize=8,
    )
ax_time.set_xticks(range(len(geom_shorts)))
ax_time.set_xticklabels(geom_shorts, fontsize=10)
ax_time.set_ylabel("Time (ms)", fontsize=11)
ax_time.set_title("Solve Time (JIT-compiled, 64x64)", fontsize=12)

# ── Residual bars ──
residuals_plot = [max(r, 1e-15) for r in residuals_list]
ax_res.bar(
    x_positions,
    residuals_plot,
    width=group_width / 3.5,
    color=bar_colors,
    edgecolor="black",
    linewidth=0.5,
)
ax_res.set_yscale("log")
ax_res.set_xticks(range(len(geom_shorts)))
ax_res.set_xticklabels(geom_shorts, fontsize=10)
ax_res.set_ylabel("Relative Residual ||r|| / ||b||", fontsize=11)
ax_res.set_title("Accuracy (Relative Residual Norm)", fontsize=12)
ax_res.set_ylim(bottom=1e-16)

# ── Legend ──
legend_entries = []
seen: set[str] = set()
for sname in solver_names:
    if sname not in seen:
        seen.add(sname)
        legend_entries.append(Patch(facecolor=colors.get(sname, "#888"), label=sname))
ax_res.legend(handles=legend_entries, loc="upper left", fontsize=9)
ax_time.legend(handles=legend_entries, loc="upper left", fontsize=9)

plt.tight_layout()
fig.savefig(IMG_DIR / "accuracy_timing.png", dpi=150, bbox_inches="tight")
plt.show()
print("Saved accuracy_timing.png")

print("\nSummary chart:") geom_order = list(results.keys()) geom_shorts = [g.split("\n")[0] for g in geom_order] colors = { "Spectral": "#2196F3", "Capacitance": "#4CAF50", "CG (bare)": "#BDBDBD", "CG": "#FF9800", "Multigrid": "#9C27B0", "MG+CG": "#E91E63", } x_positions: list[float] = [] bar_colors: list[str] = [] times_list: list[float] = [] residuals_list: list[float] = [] solver_names: list[str] = [] group_width = 0.8 for gi, geom in enumerate(geom_order): solvers = list(results[geom].keys()) n_s = len(solvers) bar_w = group_width / n_s for si, sname in enumerate(solvers): x_positions.append(gi + (si - (n_s - 1) / 2) * bar_w) bar_colors.append(colors.get(sname, "#888888")) times_list.append(results[geom][sname]["time_ms"]) residuals_list.append(results[geom][sname]["rel_residual"]) solver_names.append(sname) fig, (ax_res, ax_time) = plt.subplots(1, 2, figsize=(16, 5)) # ── Timing bars ── bars = ax_time.bar( x_positions, times_list, width=group_width / 3.5, color=bar_colors, edgecolor="black", linewidth=0.5, ) for bar, t in zip(bars, times_list, strict=True): ax_time.text( bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.15, f"{t:.1f}", ha="center", va="bottom", fontsize=8, ) ax_time.set_xticks(range(len(geom_shorts))) ax_time.set_xticklabels(geom_shorts, fontsize=10) ax_time.set_ylabel("Time (ms)", fontsize=11) ax_time.set_title("Solve Time (JIT-compiled, 64x64)", fontsize=12) # ── Residual bars ── residuals_plot = [max(r, 1e-15) for r in residuals_list] ax_res.bar( x_positions, residuals_plot, width=group_width / 3.5, color=bar_colors, edgecolor="black", linewidth=0.5, ) ax_res.set_yscale("log") ax_res.set_xticks(range(len(geom_shorts))) ax_res.set_xticklabels(geom_shorts, fontsize=10) ax_res.set_ylabel("Relative Residual ||r|| / ||b||", fontsize=11) ax_res.set_title("Accuracy (Relative Residual Norm)", fontsize=12) ax_res.set_ylim(bottom=1e-16) # ── Legend ── legend_entries = [] seen: set[str] = set() for sname in solver_names: if sname not in seen: seen.add(sname) legend_entries.append(Patch(facecolor=colors.get(sname, "#888"), label=sname)) ax_res.legend(handles=legend_entries, loc="upper left", fontsize=9) ax_time.legend(handles=legend_entries, loc="upper left", fontsize=9) plt.tight_layout() fig.savefig(IMG_DIR / "accuracy_timing.png", dpi=150, bbox_inches="tight") plt.show() print("Saved accuracy_timing.png")

8. Inhomogeneous Boundary Conditions (Known Boundary Values)¶

All solvers above assume homogeneous Dirichlet BCs ($\psi = 0$ on the boundary). In practice, boundary values are often known but non-zero — e.g., prescribed SSH from observations, reanalysis (ERA5), or a coarser model providing boundary forcing.

The standard approach is the lifting trick: decompose

$$ \psi = \psi_{\text{lift}} + \psi_{\text{hom}} $$

where $\psi_{\text{lift}}$ holds the prescribed boundary data $g$ (and is zero elsewhere), and $\psi_{\text{hom}}$ solves the corrected equation with zero BCs on the remaining cells:

$$ (A - \lambda)\,\psi_{\text{hom}} = f - (A - \lambda)\,\psi_{\text{lift}} $$

finitevolX does this for you: every convenience wrapper takes known_values. (The Known Boundary Values tutorial notebook derives the method from scratch.) The known cells are the coast ring — the dry cells adjacent to water, around the walls and any islands, exactly where method="cg" puts its $\psi = 0$ wall — and the solution equals known_values exactly there; every wet cell is solved for.

Example: prescribed sinusoidal SSH on basin walls¶

We solve the same Helmholtz problem on the basin domain, but now with non-zero Dirichlet data: $\psi = g$ on the boundary where $g = 0.1 \sin(2\pi y / L_y)$ (mimicking a prescribed SSH gradient from a parent model or observations).

In [ ]:

# --- Prescribed boundary data ---
# known_values is a full (Ny, Nx) field; only its coast cells are read.
g_field = jnp.asarray(0.1 * np.sin(2 * np.pi * Y / Ny))
ring = fvx.coast_ring(mask_basin_jnp)
print(f"Coast (known) cells: {int(ring.sum())}")

# --- Prescribed boundary data --- # known_values is a full (Ny, Nx) field; only its coast cells are read. g_field = jnp.asarray(0.1 * np.sin(2 * np.pi * Y / Ny)) ring = fvx.coast_ring(mask_basin_jnp) print(f"Coast (known) cells: {int(ring.sum())}")

In [ ]:

# --- One call per solver ---
psi_full_inhom = fvx.streamfunction_from_vorticity(
    rhs_basin, dx, dy, lambda_=lambda_,
    method="cg", mask=mask_basin_jnp, known_values=g_field,
)

# The capacitance solver is built on the SAME wet mask: its unknowns are
# every wet cell with psi = 0 on the coast, the same problem CG solves.
cap_inhom = fvx.build_capacitance_solver(
    mask_basin.astype(bool), dx, dy, lambda_=lambda_, base_bc="dst"
)
psi_full_cap = fvx.streamfunction_from_vorticity(
    rhs_basin, dx, dy, lambda_=lambda_, method="capacitance",
    mask=mask_basin_jnp, capacitance_solver=cap_inhom, known_values=g_field,
)

# Checks: boundary data reproduced exactly, PDE satisfied on the solved cells
# (the stencil reads the coast values, so apply it unmasked)
domain = fvx.SolveDomain(mask_basin_jnp)
eff = domain.effective_mask.astype(float)
A_full = lambda x: fvx.masked_laplacian(x, jnp.ones_like(x), dx, dy, lambda_=lambda_)
residual = (rhs_basin - A_full(psi_full_inhom)) * eff
print(f"Max |psi - g| on coast:      {float(jnp.max(jnp.abs((psi_full_inhom - g_field)[ring]))):.2e}")
print(f"Max residual on solved cells: {float(jnp.max(jnp.abs(residual))):.2e}")
print(f"CG vs capacitance:           {float(jnp.max(jnp.abs(psi_full_inhom - psi_full_cap))):.2e}")

# --- One call per solver --- psi_full_inhom = fvx.streamfunction_from_vorticity( rhs_basin, dx, dy, lambda_=lambda_, method="cg", mask=mask_basin_jnp, known_values=g_field, ) # The capacitance solver is built on the SAME wet mask: its unknowns are # every wet cell with psi = 0 on the coast, the same problem CG solves. cap_inhom = fvx.build_capacitance_solver( mask_basin.astype(bool), dx, dy, lambda_=lambda_, base_bc="dst" ) psi_full_cap = fvx.streamfunction_from_vorticity( rhs_basin, dx, dy, lambda_=lambda_, method="capacitance", mask=mask_basin_jnp, capacitance_solver=cap_inhom, known_values=g_field, ) # Checks: boundary data reproduced exactly, PDE satisfied on the solved cells # (the stencil reads the coast values, so apply it unmasked) domain = fvx.SolveDomain(mask_basin_jnp) eff = domain.effective_mask.astype(float) A_full = lambda x: fvx.masked_laplacian(x, jnp.ones_like(x), dx, dy, lambda_=lambda_) residual = (rhs_basin - A_full(psi_full_inhom)) * eff print(f"Max |psi - g| on coast: {float(jnp.max(jnp.abs((psi_full_inhom - g_field)[ring]))):.2e}") print(f"Max residual on solved cells: {float(jnp.max(jnp.abs(residual))):.2e}") print(f"CG vs capacitance: {float(jnp.max(jnp.abs(psi_full_inhom - psi_full_cap))):.2e}")

In [ ]:

# --- Look inside: the lift and the homogeneous correction ---
# SolveDomain / KnownValueLifting are the building blocks the wrapper uses;
# call them directly to plug the lifting into your own solver.
lifter = fvx.KnownValueLifting(domain, dx, dy, lambda_)
rhs_corrected, psi_lift = lifter.preprocess(rhs_basin, g_field)
sol_hom = (psi_full_inhom - psi_lift) * eff

fig, axes = plt.subplots(1, 3, figsize=(16, 4))
wet = mask_basin > 0.5

# (a) Prescribed boundary data on the coast ring
lift_np = np.where(np.asarray(ring), np.asarray(psi_lift), np.nan)
im0 = axes[0].imshow(lift_np, origin="lower", cmap="RdBu_r", interpolation="nearest")
axes[0].set_title("$\\psi_{\\mathrm{lift}}$ (boundary data $g$ on the ring)", fontsize=11)
fig.colorbar(im0, ax=axes[0], shrink=0.8)

# (b) Homogeneous correction (the wet cells)
sol_np = np.where(np.asarray(eff) > 0.5, np.asarray(sol_hom), np.nan)
im1 = axes[1].imshow(sol_np, origin="lower", cmap="RdBu_r", interpolation="nearest")
axes[1].set_title("$\\psi_{\\mathrm{hom}}$ (correction, $\\psi=0$ on the coast)", fontsize=11)
fig.colorbar(im1, ax=axes[1], shrink=0.8)

# (c) Full solution = lift + hom on the wet domain and its coast
full_np = np.where(wet | np.asarray(ring), np.asarray(psi_full_inhom), np.nan)
im2 = axes[2].imshow(full_np, origin="lower", cmap="RdBu_r", interpolation="nearest")
axes[2].set_title("$\\psi = \\psi_{\\mathrm{lift}} + \\psi_{\\mathrm{hom}}$", fontsize=11)
fig.colorbar(im2, ax=axes[2], shrink=0.8)

for ax in axes:
    ax.axis("off")

fig.suptitle(
    "Inhomogeneous Dirichlet BCs via lifting trick",
    fontsize=14,
    y=1.02,
)
plt.tight_layout()
fig.savefig(IMG_DIR / "inhomogeneous_bc.png", dpi=150, bbox_inches="tight")
plt.show()
print("Saved inhomogeneous_bc.png")

# --- Look inside: the lift and the homogeneous correction --- # SolveDomain / KnownValueLifting are the building blocks the wrapper uses; # call them directly to plug the lifting into your own solver. lifter = fvx.KnownValueLifting(domain, dx, dy, lambda_) rhs_corrected, psi_lift = lifter.preprocess(rhs_basin, g_field) sol_hom = (psi_full_inhom - psi_lift) * eff fig, axes = plt.subplots(1, 3, figsize=(16, 4)) wet = mask_basin > 0.5 # (a) Prescribed boundary data on the coast ring lift_np = np.where(np.asarray(ring), np.asarray(psi_lift), np.nan) im0 = axes[0].imshow(lift_np, origin="lower", cmap="RdBu_r", interpolation="nearest") axes[0].set_title("$\\psi_{\\mathrm{lift}}$ (boundary data $g$ on the ring)", fontsize=11) fig.colorbar(im0, ax=axes[0], shrink=0.8) # (b) Homogeneous correction (the wet cells) sol_np = np.where(np.asarray(eff) > 0.5, np.asarray(sol_hom), np.nan) im1 = axes[1].imshow(sol_np, origin="lower", cmap="RdBu_r", interpolation="nearest") axes[1].set_title("$\\psi_{\\mathrm{hom}}$ (correction, $\\psi=0$ on the coast)", fontsize=11) fig.colorbar(im1, ax=axes[1], shrink=0.8) # (c) Full solution = lift + hom on the wet domain and its coast full_np = np.where(wet | np.asarray(ring), np.asarray(psi_full_inhom), np.nan) im2 = axes[2].imshow(full_np, origin="lower", cmap="RdBu_r", interpolation="nearest") axes[2].set_title("$\\psi = \\psi_{\\mathrm{lift}} + \\psi_{\\mathrm{hom}}$", fontsize=11) fig.colorbar(im2, ax=axes[2], shrink=0.8) for ax in axes: ax.axis("off") fig.suptitle( "Inhomogeneous Dirichlet BCs via lifting trick", fontsize=14, y=1.02, ) plt.tight_layout() fig.savefig(IMG_DIR / "inhomogeneous_bc.png", dpi=150, bbox_inches="tight") plt.show() print("Saved inhomogeneous_bc.png")

known_values works with the ψ wrappers (streamfunction_from_vorticity, pv_inversion) and every method; pressure_from_divergence refuses it (its T-point Dirichlet wall is on the cell faces):

Method	How to call it
cg	method="cg", mask=mask, known_values=g (any mask, islands included)
capacitance	build the solver on the wet mask, then method="capacitance", known_values=g
spectral	bc="dst", known_values=g — the ψ basin of the rectangle, values on the wall nodes 0 / N-2
multigrid	method="cg" with preconditioner=make_multigrid_preconditioner(mg), mg built on SolveDomain(mask).effective_mask
custom	SolveDomain + KnownValueLifting.preprocess → your solver on domain.effective_mask → postprocess

Interior observations (tide gauges, altimetry) are pinned the same way with known_mask= (CG), pv_inversion takes per-layer known_values, and bc= also accepts a BoundaryConditionSet whose non-zero Dirichlet1D face values become known values on the walls. This covers:

• Observation-derived boundary data (tide gauges, altimetry)
• Reanalysis forcing (ERA5 SSH or currents at open boundaries)
• Nesting (parent model provides boundary values for a regional child)

9. Summary Table¶

Geometry	Solver	Recommended when	Complexity
Rectangle	Spectral (DST)	Full grid, constant coeff	$O(N \log N)$, exact
Fixed mask, many solves	Capacitance matrix	Constant coeff, setup affordable	$O(N \log N + B^2)$
Any mask	MG + CG	Tight tolerance, constant or variable coeff	$O(k \cdot N)$, $k$ independent of $N$
Any mask	Multigrid	A few V-cycles are accurate enough	$O(N)$ per V-cycle
Any mask, no setup	CG + spectral PC	Constant coeff	$O(k \cdot N \log N)$, $k$ grows with $N$

Measured timings, iteration counts and errors at 128² and 256² are in the Solver Comparison. See the Elliptic Solvers docs for full theory and the Preconditioner Guide for detailed comparisons.

In [ ]:

print("\n" + "=" * 80)
print(
    f"{'Geometry':<22} {'Solver':<14} {'Time (ms)':>10} {'Rel. Residual':>14} {'Detail'}"
)
print("=" * 80)
for geom in geom_order:
    geom_short = geom.split("\n")[0]
    for sname, data in results[geom].items():
        print(
            f"{geom_short:<22} {sname:<14} {data['time_ms']:>10.2f} "
            f"{data['rel_residual']:>14.1e} {data['label']}"
        )
    print("-" * 80)

print("\n" + "=" * 80) print( f"{'Geometry':<22} {'Solver':<14} {'Time (ms)':>10} {'Rel. Residual':>14} {'Detail'}" ) print("=" * 80) for geom in geom_order: geom_short = geom.split("\n")[0] for sname, data in results[geom].items(): print( f"{geom_short:<22} {sname:<14} {data['time_ms']:>10.2f} " f"{data['rel_residual']:>14.1e} {data['label']}" ) print("-" * 80)