Tutorial: matrix-free anisotropic PDE with a line preconditioner¶

This tutorial solves a two-dimensional anisotropic diffusion-reaction equation without passing an assembled matrix to FGMRES:

\[ -\epsilon_x u_{xx}-\epsilon_yu_{yy}+cu=f, \qquad \epsilon_x\gg\epsilon_y. \]

Point Jacobi treats the diagonal but not the strong \(x\) coupling. We construct an exact tridiagonal solve along every \(x\) line and use it as a preconditioner.

1. Matrix-free stencil¶

import jax
import jax.numpy as jnp
import solvax as sx

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

nx, ny = 128, 32
eps_x, eps_y, reaction = 1.0, 1e-2, 0.5
hx, hy = 1.0 / (nx + 1), 1.0 / (ny + 1)
wx, wy = eps_x / hx**2, eps_y / hy**2
diagonal = 2.0 * wx + 2.0 * wy + reaction

def matvec(flat):
    u = flat.reshape(nx, ny)
    padded = jnp.pad(u, ((1, 1), (1, 1)))
    center = diagonal * u
    x_part = -wx * (padded[:-2, 1:-1] + padded[2:, 1:-1])
    y_part = -wy * (padded[1:-1, :-2] + padded[1:-1, 2:])
    return (center + x_part + y_part).reshape(-1)

The padding encodes homogeneous Dirichlet boundaries. The function accepts and returns (nx * ny,), matching FGMRES.

2. Exact solves along the stiff direction¶

lower_x = jnp.full((nx, ny), -wx)
diag_x = jnp.full((nx, ny), diagonal)
upper_x = jnp.full((nx, ny), -wx)

def x_line_inverse(residual_flat):
    residual = residual_flat.reshape(nx, ny)
    correction = sx.tridiagonal_solve(
        lower_x, diag_x, upper_x, residual, method="auto"
    )
    return correction.reshape(-1)

The leading nx axis is the tridiagonal system; all ny columns are solved in one call.

3. Build the smoother/preconditioner¶

preconditioner = sx.line_smoother(
    matvec,
    [x_line_inverse],
    omega=0.9,
    sweeps=2,
)

Each sweep computes a true outer residual and applies the line inverse. Because the preconditioner is an iterative action, flexible GMRES is appropriate.

4. Solve and inspect diagnostics¶

rhs = jnp.ones(nx * ny)
solution = sx.gmres(
    matvec,
    rhs,
    precond=preconditioner,
    restart=40,
    rtol=1e-9,
    atol=1e-12,
    max_restarts=30,
)

if not solution.converged:
    raise RuntimeError(
        f"FGMRES stopped after {int(solution.iterations)} steps; "
        f"residual={float(solution.residual_norm):.3e}"
    )

5. Compare against point Jacobi¶

point = sx.jacobi(jnp.full(nx * ny, diagonal))
point_solution = sx.gmres(
    matvec, rhs, precond=point, restart=40, rtol=1e-9, max_restarts=30
)

print("point iterations:", int(point_solution.iterations))
print("line iterations:", int(solution.iterations))

The meaningful comparison is total time including the preconditioner, not only iteration count. On a GPU, the fused tridiagonal backend can make the line solve especially attractive.

6. Extensions¶

Add a \(y\)-line inverse and alternate directions when anisotropy changes spatially or both directions are stiff.
Place line smoothing inside p_multigrid to eliminate high-frequency error while a coarse grid corrects smooth error.
Use PCG only if the complete preconditioner action is symmetric positive definite. The sequential relaxed line smoother above is safest with FGMRES.
Replace the stencil with spectral or finite-volume code; the outer API is unchanged as long as matvec preserves the flat shape.