Chapter 11: Scientific Computing and Differentiable Simulation

# First thing: turn on 64-bit precision (must happen before any JAX array is created)
import jax
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp

# Self-check: the default float type, random numbers and Python scalar conversion are all float64
assert jnp.zeros(3).dtype == jnp.float64
assert jnp.asarray(1.0).dtype == jnp.float64
assert jax.random.normal(jax.random.key(0), (2,)).dtype == jnp.float64
assert jnp.arange(3).dtype == jnp.int64
print("Default float dtype:", jnp.zeros(3).dtype)

# Common imports
import os, sys, time, json, inspect, subprocess, textwrap
from functools import partial

import numpy as np
import scipy
import scipy.linalg
import scipy.sparse as sps
import scipy.sparse.linalg as spsl
import optax
import matplotlib
import matplotlib.pyplot as plt
from jax import lax
import jax.scipy.sparse.linalg as jsl
from jax.experimental import sparse
from jax.ad_checkpoint import print_saved_residuals

plt.rcParams["font.sans-serif"] = ["PingFang SC", "Heiti SC", "Arial Unicode MS", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False

print("JAX", jax.__version__, "| optax", optax.__version__, "| scipy", scipy.__version__,
      "| numpy", np.__version__, "| matplotlib", matplotlib.__version__)
print("Devices:", jax.devices())

def rel(a, b):
    """Relative error ||a-b|| / ||b|| (computed after converting everything to NumPy float64)"""
    a, b = np.asarray(a, np.float64), np.asarray(b, np.float64)
    return float(np.linalg.norm(a - b) / np.linalg.norm(b))

def timeit(fn, *args, n=20):
    """Warm up (compile) first, then time n runs and return the mean in milliseconds. Only for printing, never asserted."""
    jax.block_until_ready(fn(*args))
    t = time.perf_counter()
    for _ in range(n):
        jax.block_until_ready(fn(*args))
    return (time.perf_counter() - t) / n * 1e3

# ❌ Wrong: create arrays first, turn on x64 later; the existing arrays stay float32 and spread through weakly typed operations
probe = textwrap.dedent("""
    import jax, jax.numpy as jnp, json
    a = jnp.ones(3)                                  # the default is still float32 at this point
    jax.config.update("jax_enable_x64", True)        # too late
    b = jnp.ones(3)
    print(json.dumps({"a": str(a.dtype), "b": str(b.dtype), "a*1.0": str((a * 1.0).dtype),
                      "a+b": str((a + b).dtype)}))
""")
try:
    r = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, timeout=120,
                       env={k: v for k, v in os.environ.items() if k != "JAX_ENABLE_X64"})
    late = json.loads(r.stdout.strip().splitlines()[-1])
    print("Measured dtypes:", late)
except Exception as e:
    late = None
    print("Subprocess failed:", type(e).__name__, e)
assert late["a"] == "float32" and late["a*1.0"] == "float32"   # the old array is still single precision, and so is the result of multiplying by a scalar
assert late["b"] == "float64"

# ✅ Correct: update first, then create any arrays (also checked in a subprocess; section 0 of this kernel does exactly this)
probe_ok = textwrap.dedent("""
    import jax
    jax.config.update("jax_enable_x64", True)        # first thing
    import jax.numpy as jnp, json
    a = jnp.ones(3)
    print(json.dumps({"a": str(a.dtype), "a*1.0": str((a * 1.0).dtype)}))
""")
r = subprocess.run([sys.executable, "-c", probe_ok], capture_output=True, text=True, timeout=120,
                   env={k: v for k, v in os.environ.items() if k != "JAX_ENABLE_X64"})
early = json.loads(r.stdout.strip().splitlines()[-1])
print("Measured dtypes:", early)
assert early == {"a": "float64", "a*1.0": "float64"}
assert jnp.ones(3).dtype == jnp.float64 and (jnp.ones(3) * 1.0).dtype == jnp.float64   # this kernel

# ❌ Wrong: doing a Poisson convergence study in float32; on fine meshes rounding error swamps discretization error and second-order convergence "disappears"
probe32 = textwrap.dedent("""
    import jax, jax.numpy as jnp, json
    import jax.scipy.sparse.linalg as jsl
    assert jnp.zeros(1).dtype == jnp.float32
    def neg_lap(u, h):
        up = jnp.pad(u, 1)
        return (4 * u - up[:-2, 1:-1] - up[2:, 1:-1] - up[1:-1, :-2] - up[1:-1, 2:]) / h**2
    solve = jax.jit(lambda f, h: jsl.cg(lambda v: neg_lap(v, h), f, tol=1e-10, maxiter=4 * f.size)[0])
    errs = []
    for N in [15, 31, 63, 127, 255]:
        h = 1.0 / (N + 1); x = jnp.linspace(h, 1 - h, N); X, Y = jnp.meshgrid(x, x, indexing="ij")
        ue = jnp.sin(jnp.pi * X) * jnp.sin(jnp.pi * Y)
        errs.append(float(jnp.max(jnp.abs(solve(2 * jnp.pi**2 * ue, h) - ue))))
    fr = jax.random.normal(jax.random.key(0), (63, 63))           # random right-hand side (not an eigenvector)
    res63 = float(jnp.linalg.norm(neg_lap(solve(fr, 1 / 64), 1 / 64) - fr) / jnp.linalg.norm(fr))
    g = lambda a: jnp.sum(jnp.sin(a * jnp.arange(10.0)))
    a, eps = jnp.float32(0.7), 1e-6
    print(json.dumps({"errs": errs, "fd": float((g(a + eps) - g(a - eps)) / (2 * eps)), "ad": float(jax.grad(g)(a)),
                      "fd_best": float((g(a + 1e-3) - g(a - 1e-3)) / 2e-3), "res63": res63}))
""")
try:
    r = subprocess.run([sys.executable, "-c", probe32], capture_output=True, text=True, timeout=300,
                       env={k: v for k, v in os.environ.items() if k != "JAX_ENABLE_X64"})
    res32 = json.loads(r.stdout.strip().splitlines()[-1])
    e32 = np.array(res32["errs"])
    print("float32 max error:", np.array2string(e32, precision=3))
    print("Error ratio between consecutive meshes:", np.round(e32[:-1] / e32[1:], 3), "  (should be about 4 for second-order convergence)")
    print(f"float32 finite-difference gradient {res32['fd']:.6f}  vs  AD gradient {res32['ad']:.6f}"
          f"  relative error {abs(res32['fd'] - res32['ad']) / abs(res32['ad']):.1e}")
    print(f"float32 finite difference (step 1e-3, near optimal) relative error {abs(res32['fd_best'] - res32['ad']) / abs(res32['ad']):.1e}"
          f"; true relative residual of CG (tol=1e-10, N=63, random right-hand side) {res32['res63']:.1e}")
except Exception as e:
    res32 = None
    print("Subprocess failed:", type(e).__name__, e)
assert e32[-2] / e32[-1] < 3.5                                     # the error ratio at the finest level is clearly off from 4
assert abs(res32["fd"] - res32["ad"]) / abs(res32["ad"]) > 1e-3    # with step 1e-6, float32 finite differences keep only 2 to 3 significant digits
assert 1e-6 < abs(res32["fd_best"] - res32["ad"]) / abs(res32["ad"]) < 1e-4   # even near the optimal step, only 4 to 5 digits
assert res32["res63"] > 1e-7                                       # CG's true residual can't get below 1e-7

# ✅ Correct: gradient check in float64 (for the convergence study see section 2; second-order convergence holds all the way to the finest mesh)
g = lambda a: jnp.sum(jnp.sin(a * jnp.arange(10.0)))
a0, eps = 0.7, 1e-6
fd64 = (g(a0 + eps) - g(a0 - eps)) / (2 * eps)
ad64 = jax.grad(g)(a0)
print(f"float64 finite difference {float(fd64):.10f}  vs  AD {float(ad64):.10f}  relative error {abs(fd64 - ad64) / abs(ad64):.1e}")
assert abs(fd64 - ad64) / abs(ad64) < 1e-8
# For the Poisson convergence order in float64, see section 2

def neg_lap(u, h):
    """Matrix-free five-point operator -Δ_h u (homogeneous Dirichlet boundary, implemented by zero padding)"""
    up = jnp.pad(u, 1)
    return (4 * u - up[:-2, 1:-1] - up[2:, 1:-1] - up[1:-1, :-2] - up[1:-1, 2:]) / h**2

def scipy_neg_lap(N, h):
    """SciPy reference: assemble the same matrix with Kronecker products (CSR, float64)"""
    T = sps.diags([-np.ones(N - 1), 2 * np.ones(N), -np.ones(N - 1)], [-1, 0, 1])
    I = sps.identity(N)
    return ((sps.kron(I, T) + sps.kron(T, I)) / h**2).tocsr()

N = 31; h = 1.0 / (N + 1)
v_np = np.random.default_rng(0).standard_normal((N, N))
A_ref = scipy_neg_lap(N, h)
ref = (A_ref @ v_np.ravel()).reshape(N, N)            # row-major flattening, matching the ordering of kron(I, T) + kron(T, I)
err_pad = rel(neg_lap(jnp.asarray(v_np), h), ref)
print(f"pad-and-slice operator vs SciPy CSR: relative error {err_pad:.1e}")
assert err_pad < 1e-14

# ❌ Wrong: taking neighbors directly with jnp.roll; roll wraps values from the opposite edge around, which amounts to periodic boundary conditions
def neg_lap_roll_wrong(u, h):
    return (4 * u - jnp.roll(u, 1, 0) - jnp.roll(u, -1, 0) - jnp.roll(u, 1, 1) - jnp.roll(u, -1, 1)) / h**2

try:
    err_roll = rel(neg_lap_roll_wrong(jnp.asarray(v_np), h), ref)
    out_roll = np.asarray(neg_lap_roll_wrong(jnp.asarray(v_np), h))
    err_inner = rel(out_roll[1:-1, 1:-1], ref[1:-1, 1:-1])
    err_edge = rel(out_roll[0], ref[0])
    print(f"plain-roll operator vs SciPy (Dirichlet): overall relative error {err_roll:.3f}; "
          f"interior nodes {err_inner:.1e}, first row {err_edge:.3f}; no error raised, but every node next to the boundary is wrong")
except Exception as e:
    err_roll = None
    print(type(e).__name__, e)
assert err_roll > 0.05 and err_edge > 0.1 and err_inner < 1e-14

# ✅ Correct: pad a ring of zeros first, roll on the padded array, then crop the outer ring (exactly equivalent to the slicing version)
def neg_lap_roll(u, h):
    up = jnp.pad(u, 1)
    r = (4 * up - jnp.roll(up, 1, 0) - jnp.roll(up, -1, 0) - jnp.roll(up, 1, 1) - jnp.roll(up, -1, 1)) / h**2
    return r[1:-1, 1:-1]

err_roll_ok = rel(neg_lap_roll(jnp.asarray(v_np), h), ref)
print(f"pad + roll operator vs SciPy: relative error {err_roll_ok:.1e}")
assert err_roll_ok < 1e-14
assert jnp.allclose(neg_lap_roll(jnp.asarray(v_np), h), neg_lap(jnp.asarray(v_np), h), rtol=0, atol=1e-9)

@jax.jit
def solve_poisson(f, h):
    """-Δu = f, homogeneous Dirichlet; the operator is passed to cg as a function"""
    u, info = jsl.cg(lambda v: neg_lap(v, h), f, tol=1e-12, maxiter=10 * f.size)
    return u

def u_star(p):
    """Manufactured solution u*(x, y), zero on the boundary of the unit square"""
    x, y = p
    return jnp.exp(x + 0.5 * y) * jnp.sin(jnp.pi * x) * jnp.sin(jnp.pi * y)

f_star = lambda p: -jnp.trace(jax.hessian(u_star)(p))       # f = -Δu*, computed automatically by AD

# Spot-check one point against the hand-derived analytic expression
p0 = jnp.array([0.3, 0.7])
x0, y0 = p0
E, sx, cx, sy, cy = jnp.exp(x0 + 0.5 * y0), jnp.sin(jnp.pi * x0), jnp.cos(jnp.pi * x0), jnp.sin(jnp.pi * y0), jnp.cos(jnp.pi * y0)
lap_hand = E * ((1 - jnp.pi**2) * sx * sy + 2 * jnp.pi * cx * sy + (0.25 - jnp.pi**2) * sx * sy + jnp.pi * sx * cy)
assert abs(f_star(p0) + lap_hand) < 1e-12

def manufactured(N):
    """Return h, the grid coordinates, the manufactured solution u* and the source term f = -Δu* (all on the N×N interior nodes)"""
    h = 1.0 / (N + 1)
    x = jnp.linspace(h, 1 - h, N)
    X, Y = jnp.meshgrid(x, x, indexing="ij")
    P = jnp.stack([X, Y], -1).reshape(-1, 2)
    ue = jax.vmap(u_star)(P).reshape(N, N)
    f = jax.vmap(f_star)(P).reshape(N, N)
    return h, X, Y, ue, f

Ns = [15, 31, 63, 127, 255]
errs = []
print(f"{'N':>5} {'h':>9} {'max error':>12} {'rel resid':>10} {'time (ms)':>9}")
for N in Ns:
    h, X, Y, ue, f = manufactured(N)
    u = solve_poisson(f, h)
    e = float(jnp.max(jnp.abs(u - ue)))
    res_rel = rel(neg_lap(u, h), f)
    errs.append(e)
    print(f"{N:>5} {h:>9.5f} {e:>12.4e} {res_rel:>10.1e} {timeit(solve_poisson, f, h, n=3):>9.1f}")
    assert res_rel < 1e-10                                # algebraic error is far below discretization error

errs = np.array(errs)
ratios = errs[:-1] / errs[1:]
orders = np.log2(ratios)
print("Error ratios:", np.round(ratios, 4))
print("Convergence orders:", np.round(orders, 4))
assert np.all(np.abs(ratios - 4) < 0.1)                   # refining the mesh by 2 shrinks the error to about 1/4
assert np.all(np.abs(orders - 2) < 0.02)                  # second-order convergence
POISSON_ORDERS = orders

# Cross-check against SciPy's sparse direct solver (same discrete system; should agree within rounding error)
N = 63
h, X, Y, ue, f = manufactured(N)
u_jax = solve_poisson(f, h)
u_sp = spsl.spsolve(scipy_neg_lap(N, h).tocsc(), np.asarray(f).ravel()).reshape(N, N)
print(f"JAX cg vs SciPy spsolve: relative difference {rel(u_jax, u_sp):.1e}")
assert rel(u_jax, u_sp) < 1e-10

# Check: info is None; with too small a maxiter it silently returns an unconverged solution
h = 1.0 / 64
f = jnp.ones((63, 63))
u_bad, info = jsl.cg(lambda v: neg_lap(v, h), f, tol=1e-12, maxiter=5)
print("info =", info, "| relative residual after only 5 iterations:", f"{rel(neg_lap(u_bad, h), f):.2e}")
assert info is None
assert rel(neg_lap(u_bad, h), f) > 1e-2

def five_point_triplets(N, h):
    """COO triplets (row, col, val) of the five-point matrix, generated with NumPy (preprocessing doesn't need to be differentiable)"""
    idx = np.arange(N * N).reshape(N, N)
    I, J = np.meshgrid(np.arange(N), np.arange(N), indexing="ij")
    rows, cols, vals = [idx.ravel()], [idx.ravel()], [np.full(N * N, 4 / h**2)]
    for di, dj in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
        I2, J2 = I + di, J + dj
        m = (I2 >= 0) & (I2 < N) & (J2 >= 0) & (J2 < N)
        rows.append(idx[I[m], J[m]]); cols.append(idx[I2[m], J2[m]]); vals.append(np.full(m.sum(), -1 / h**2))
    return np.concatenate(rows), np.concatenate(cols), np.concatenate(vals)

N = 63; h = 1.0 / (N + 1)
r_, c_, v_ = five_point_triplets(N, h)
A_csr = sps.csr_matrix((v_, (r_, c_)), shape=(N * N, N * N))                       # SciPy reference
A_bcoo = sparse.BCOO((jnp.asarray(v_), jnp.stack([jnp.asarray(r_), jnp.asarray(c_)], axis=1)),
                     shape=(N * N, N * N))
print(A_bcoo, "| data:", A_bcoo.data.dtype, "| indices:", A_bcoo.indices.dtype)

x_np = np.random.default_rng(1).standard_normal(N * N)
matvec = jax.jit(lambda A, x: A @ x)                    # A is passed in as an argument, not captured in a closure
y_bcoo = matvec(A_bcoo, jnp.asarray(x_np))
y_csr = A_csr @ x_np
print(f"BCOO matvec vs SciPy CSR: relative error {rel(y_bcoo, y_csr):.1e}")
assert rel(y_bcoo, y_csr) < 1e-14
assert rel(neg_lap(jnp.asarray(x_np).reshape(N, N), h).ravel(), y_csr) < 1e-14   # matches the matrix-free operator

# Convert directly from SciPy, and convert back to dense for a small-scale check
A_from_sp = sparse.BCOO.from_scipy_sparse(A_csr)
assert A_from_sp.nse == A_csr.nnz
assert np.allclose(sparse.BCOO.from_scipy_sparse(scipy_neg_lap(5, 1 / 6)).todense(),
                   scipy_neg_lap(5, 1 / 6).toarray())

# Finite-element-style assembly: duplicate indices accumulate (BCOO keeps duplicates, sums them in matvec; sum_duplicates merges them)
A_dup = sparse.BCOO((jnp.array([1.0, 2.0, 3.0, 4.0]),
                     jnp.array([[0, 0], [0, 0], [1, 1], [0, 1]])), shape=(2, 2))
print("Dense form:\n", A_dup.todense(), "\nnse:", A_dup.nse, "→ nse after merging:", A_dup.sum_duplicates().nse)
assert np.allclose(A_dup.todense(), [[3.0, 4.0], [0.0, 3.0]])
assert A_dup.sum_duplicates().nse == 3
assert np.allclose(A_dup @ jnp.array([1.0, 1.0]), [7.0, 3.0])

def energy_dense(M, x):
    """Written as for a dense matrix: ||Mx||^2 + x^T M^T x"""
    return jnp.sum((M @ x) ** 2) + (M.T @ x) @ x

energy_sparse = jax.jit(sparse.sparsify(energy_dense))
xj = jnp.asarray(x_np)
e_sp = energy_sparse(A_bcoo, xj)
e_ref = np.sum((A_csr @ x_np) ** 2) + (A_csr.T @ x_np) @ x_np
print(f"sparsify result {float(e_sp):.6e}  vs  SciPy {e_ref:.6e}")
assert abs(float(e_sp) - e_ref) / abs(e_ref) < 1e-13

# Derivative with respect to the matrix nonzeros: d/dA_ij ||Ax||^2 = 2 (Ax)_i x_j, evaluated only at the nonzero positions
def loss_data(data, indices, x):
    M = sparse.BCOO((data, indices), shape=(N * N, N * N))
    return jnp.sum((M @ x) ** 2)

g_data = jax.jit(jax.grad(loss_data))(A_bcoo.data, A_bcoo.indices, xj)
g_ref = 2 * (A_csr @ x_np)[r_] * x_np[c_]
print("Shape of the gradient w.r.t. data:", g_data.shape, "| relative error vs the analytic formula:", f"{rel(g_data, g_ref):.1e}")
assert rel(g_data, g_ref) < 1e-13

g_sp = sparse.grad(lambda M, x: jnp.sum((M @ x) ** 2))(A_bcoo, xj)
assert isinstance(g_sp, sparse.BCOO) and g_sp.nse == A_bcoo.nse
assert rel(g_sp.data, g_ref) < 1e-13

from jax.experimental.sparse.linalg import spsolve

b = jnp.ones(N * N)
u_cg_bcoo = jax.jit(lambda A, b: jsl.cg(lambda v: A @ v, b, tol=1e-12, maxiter=10 * b.size)[0])(A_bcoo, b)
u_ref = spsl.spsolve(A_csr.tocsc(), np.ones(N * N))
print(f"BCOO + cg vs SciPy spsolve: relative difference {rel(u_cg_bcoo, u_ref):.1e}")
assert rel(u_cg_bcoo, u_ref) < 1e-10

A_bcsr = sparse.BCSR.from_bcoo(A_bcoo)                  # convert to the CSR triple
u_ss = jax.jit(spsolve)(A_bcsr.data, A_bcsr.indices, A_bcsr.indptr, b)
g_ss = jax.grad(lambda b: spsolve(A_bcsr.data, A_bcsr.indices, A_bcsr.indptr, b).sum())(b)
print(f"spsolve (CPU calls back into SciPy) vs SciPy: {rel(u_ss, u_ref):.1e}; works with jit and grad")
assert rel(u_ss, u_ref) < 1e-12
assert rel(g_ss, spsl.spsolve(A_csr.T.tocsc(), np.ones(N * N))) < 1e-12   # gradient = A^{-T} 1
assert "scipy.sparse.linalg.spsolve" in inspect.getsource(sparse.linalg._spsolve_cpu_lowering)

# ❌ Wrong: vmap over spsolve (solving for several right-hand sides at once); there's no batching rule
B = jnp.stack([b, 2 * b, 3 * b])
try:
    jax.vmap(lambda bb: spsolve(A_bcsr.data, A_bcsr.indices, A_bcsr.indptr, bb))(B)
    vmap_err = None
except NotImplementedError as e:
    vmap_err = str(e)
    print("NotImplementedError:", vmap_err)
assert vmap_err is not None and "spsolve" in vmap_err

# ✅ Correct: for batched right-hand sides use vmap(cg) (the matvec is a BCOO, which batches); or factor once in SciPy and back-substitute many times
U = jax.jit(jax.vmap(lambda A, bb: jsl.cg(lambda v: A @ v, bb, tol=1e-12, maxiter=10 * bb.size)[0],
                     in_axes=(None, 0)))(A_bcoo, B)
lu = spsl.splu(A_csr.tocsc())                           # SciPy: one LU, many back-substitutions
U_ref = np.stack([lu.solve(np.asarray(bb)) for bb in B])
print(f"batched solve with vmap(cg) vs SciPy splu: {rel(U, U_ref):.1e}")
assert rel(U, U_ref) < 1e-10

# Timing (printed only): the same N=63 five-point operator, three kinds of matvec
u2 = xj.reshape(N, N)
t_bcoo = timeit(matvec, A_bcoo, xj, n=200)
t_mf = timeit(jax.jit(neg_lap), u2, h, n=200)
t0 = time.perf_counter()
for _ in range(200):
    A_csr @ x_np
t_sp = (time.perf_counter() - t0) / 200 * 1e3
print(f"BCOO matvec        : {t_bcoo * 1e3:8.1f} µs")
print(f"matrix-free matvec : {t_mf * 1e3:8.1f} µs")
print(f"SciPy CSR          : {t_sp * 1e3:8.1f} µs   (NumPy side, no JAX dispatch overhead)")

def apply_A(k, u, h):
    """Matrix-free variable-coefficient diffusion operator -∇·(k∇u); k and u are both N×N node arrays, and k outside the boundary takes the nearest value"""
    kp = jnp.pad(k, 1, mode="edge")
    up = jnp.pad(u, 1)                                    # Dirichlet 0
    kx = 0.5 * (kp[1:, 1:-1] + kp[:-1, 1:-1])            # face coefficients in x (N+1, N)
    ky = 0.5 * (kp[1:-1, 1:] + kp[1:-1, :-1])            # face coefficients in y (N, N+1)
    fx = kx * (up[1:, 1:-1] - up[:-1, 1:-1])              # face fluxes
    fy = ky * (up[1:-1, 1:] - up[1:-1, :-1])
    return -((fx[1:] - fx[:-1]) + (fy[:, 1:] - fy[:, :-1])) / h**2

N = 31; h = 1.0 / (N + 1)
xg = jnp.linspace(h, 1 - h, N)
XX, YY = jnp.meshgrid(xg, xg, indexing="ij")
k0 = 1 + 0.5 * jnp.exp(-((XX - 0.6) ** 2 + (YY - 0.4) ** 2) / 0.02)
ka, kb = jax.random.split(jax.random.key(10))
uu, vv = jax.random.normal(ka, (N, N)), jax.random.normal(kb, (N, N))

assert jnp.allclose(apply_A(jnp.ones((N, N)), uu, h), neg_lap(uu, h), rtol=1e-13, atol=1e-9)
lhs, rhs_ = jnp.vdot(apply_A(k0, uu, h), vv), jnp.vdot(uu, apply_A(k0, vv, h))
print(f"<Au,v> = {float(lhs):.10e}, <u,Av> = {float(rhs_):.10e}")
assert abs(lhs - rhs_) / abs(lhs) < 1e-12
A_T = jax.linear_transpose(lambda v: apply_A(k0, v, h), uu)
assert jnp.allclose(A_T(vv)[0], apply_A(k0, vv, h), rtol=1e-12, atol=1e-8)
assert jnp.vdot(uu, apply_A(k0, uu, h)) > 0              # spot-check positive definiteness
print("Variable-coefficient operator: reduction check, symmetry and linear_transpose consistency all pass")

def cg_unrolled(matvec, b, n_iter):
    """CG with a fixed n_iter steps (static-length fori_loop → reverse-mode differentiable, but stores every step)"""
    rs0 = jnp.vdot(b, b)
    def body(i, c):
        x, r, p, rs = c
        Ap = matvec(p)
        alpha = rs / jnp.vdot(p, Ap)
        x, r = x + alpha * p, r - alpha * Ap
        rs_new = jnp.vdot(r, r)
        return x, r, r + (rs_new / rs) * p, rs_new
    return lax.fori_loop(0, n_iter, body, (jnp.zeros_like(b), b, b, rs0))[0]

def cg_while(matvec, b, tol=1e-12, maxiter=10_000):
    """CG that stops on the relative residual (while_loop; can't be differentiated in reverse mode directly)"""
    bb = jnp.vdot(b, b)
    def cond(c):
        x, r, p, rs, i = c
        return (rs > tol**2 * bb) & (i < maxiter)
    def body(c):
        x, r, p, rs, i = c
        Ap = matvec(p)
        alpha = rs / jnp.vdot(p, Ap)
        x, r = x + alpha * p, r - alpha * Ap
        rs_new = jnp.vdot(r, r)
        return x, r, r + (rs_new / rs) * p, rs_new, i + 1
    return lax.while_loop(cond, body, (jnp.zeros_like(b), b, b, bb, 0))[0]

def u_unrolled(k, f, h, n_iter):
    return cg_unrolled(lambda v: apply_A(k, v, h), f, n_iter)

def u_implicit(k, f, h, tol=1e-12, maxiter=10_000):
    """Implicit differentiation: AD only uses matvec (differentiating w.r.t. the k in its closure) and solve itself"""
    solve = lambda mv, rhs: cg_while(mv, rhs, tol, maxiter)
    return lax.custom_linear_solve(lambda v: apply_A(k, v, h), f, solve, symmetric=True)

def J_unrolled(k, f, u_obs, h, n_iter):
    return 0.5 * h**2 * jnp.sum((u_unrolled(k, f, h, n_iter) - u_obs) ** 2)

def J_implicit(k, f, u_obs, h, tol=1e-12, maxiter=10_000):
    return 0.5 * h**2 * jnp.sum((u_implicit(k, f, h, tol, maxiter) - u_obs) ** 2)

f_src = jnp.ones((N, N))
u_obs = 0.9 * u_implicit(jnp.ones((N, N)), f_src, h)       # a synthetic "observation"

grad_unr = jax.jit(jax.grad(J_unrolled), static_argnums=4)
grad_imp = jax.jit(jax.grad(J_implicit))
K_UNROLL = 300
g_unr = grad_unr(k0, f_src, u_obs, h, K_UNROLL)
g_imp = grad_imp(k0, f_src, u_obs, h)

# Finite differences (central, along 3 random directions); the forward pass uses cg_while with a tight tolerance
J_fd = jax.jit(J_implicit)
eps = 1e-5
print(f"{'dir':>4} {'finite diff':>16} {'implicit':>16} {'unrolled':>16} {'impl rel err':>12} {'unr rel err':>12}")
for i, key in enumerate(jax.random.split(jax.random.key(3), 3)):
    d = jax.random.normal(key, (N, N))
    fd = (J_fd(k0 + eps * d, f_src, u_obs, h) - J_fd(k0 - eps * d, f_src, u_obs, h)) / (2 * eps)
    gi, gu = jnp.vdot(g_imp, d), jnp.vdot(g_unr, d)
    ei, eu = abs(gi - fd) / abs(fd), abs(gu - fd) / abs(fd)
    print(f"{i:>4} {float(fd):>16.10e} {float(gi):>16.10e} {float(gu):>16.10e} {float(ei):>12.1e} {float(eu):>12.1e}")
    assert ei < 1e-7                                       # the implicit gradient matches finite differences
    assert eu < 1e-3                                       # the gradient unrolled for 300 steps also matches finite differences (but several orders of magnitude worse)

GRAD_REL_UNR_IMP = rel(g_unr, g_imp)
print(f"\nUnrolled {K_UNROLL} steps vs implicit: relative gradient difference {GRAD_REL_UNR_IMP:.2e}")
assert GRAD_REL_UNR_IMP < 1e-4

# Memory: read the temporary buffers the compiled XLA program needs for backprop
def temp_mib(fn, *args, static_argnums=()):
    c = jax.jit(jax.grad(fn), static_argnums=static_argnums).lower(*args).compile()
    return c.memory_analysis().temp_size_in_bytes / 2**20

MEM = {}
for K in [100, 300]:
    MEM[f"unrolled {K} steps"] = temp_mib(J_unrolled, k0, f_src, u_obs, h, K, static_argnums=4)
MEM["implicit (custom_linear_solve)"] = temp_mib(J_implicit, k0, f_src, u_obs, h)
for name, m in MEM.items():
    print(f"{name:<28}: {m:8.3f} MiB")
assert MEM["unrolled 300 steps"] > 2.5 * MEM["unrolled 100 steps"]          # grows linearly with the number of iterations
assert MEM["implicit (custom_linear_solve)"] < MEM["unrolled 300 steps"] / 50

print("\nTiming (printed only):")
print(f"  unrolled 300 steps grad: {timeit(grad_unr, k0, f_src, u_obs, h, K_UNROLL):.2f} ms")
print(f"  implicit grad          : {timeit(grad_imp, k0, f_src, u_obs, h):.2f} ms")

print("\nResiduals saved by the implicit version (only the input k and a few same-sized arrays computed from the solution u; nothing stacked per iteration):")
print_saved_residuals(lambda k: J_implicit(k, f_src, u_obs, h), k0)

print(f"{'K':>5} {'fwd sol err':>12} {'unrolled grad err':>14} {'implicit grad err':>14}")
scan_rows = []
for K in [20, 40, 60, 70, 80, 90, 100, 150, 300]:
    gu = grad_unr(k0, f_src, u_obs, h, K)
    gi = grad_imp(k0, f_src, u_obs, h, 0.0, K)            # tol=0: exactly K iterations
    ef = rel(u_unrolled(k0, f_src, h, K), u_implicit(k0, f_src, h))
    eu, ei = rel(gu, g_imp), rel(gi, g_imp)
    scan_rows.append((K, ef, eu, ei))
    print(f"{K:>5} {ef:>12.2e} {eu:>14.2e} {ei:>14.2e}")
scan_rows = np.array(scan_rows)

# Same K=80: finite differences of the unrolled function itself give the true gradient; but both the AD jvp (forward) and vjp (reverse) are far from it, and disagree with each other
d = jax.random.normal(jax.random.key(3), (N, N))
_, jvp80 = jax.jvp(lambda k: J_unrolled(k, f_src, u_obs, h, 80), (k0,), (d,))
vjp80 = jnp.vdot(grad_unr(k0, f_src, u_obs, h, 80), d)
fd80 = (J_unrolled(k0 + 1e-6 * d, f_src, u_obs, h, 80) - J_unrolled(k0 - 1e-6 * d, f_src, u_obs, h, 80)) / 2e-6
true80 = jnp.vdot(g_imp, d)
print(f"\nK=80: true gradient {float(true80):.6e} | finite difference of the unrolled function itself {float(fd80):.6e} | "
      f"unrolled jvp {float(jvp80):.6e} | unrolled vjp {float(vjp80):.6e}")
assert abs(fd80 - true80) < 1e-6 * abs(fd80)               # the function value itself has converged, so its finite difference is the true gradient
assert abs(jvp80 - true80) > 100 * abs(true80)             # the forward-mode unrolled derivative is just as wrong → not a reverse-mode bug
assert abs(vjp80 - true80) > 100 * abs(true80)

assert scan_rows[:, 2].max() > 1e-2                          # the unrolled gradient gets badly distorted around convergence
assert np.all(scan_rows[scan_rows[:, 0] >= 80][:, 3] < 1e-5)  # the implicit gradient converges steadily along with the forward solve

n1 = 50; h1 = 1.0 / (n1 + 1)

def conv_diff(beta, u):
    """1D convection-diffusion -u'' + beta u' (central differences, Dirichlet 0); nonsymmetric for beta≠0"""
    up = jnp.pad(u, 1)
    return (2 * u - up[:-2] - up[2:]) / h1**2 + beta * (up[2:] - up[:-2]) / (2 * h1)

gmres_solve = lambda mv, rhs: jsl.gmres(mv, rhs, tol=1e-13, atol=0.0, restart=n1, maxiter=20)[0]
f1 = jnp.sin(jnp.linspace(0, 3, n1))
w1 = jnp.linspace(1, 2, n1)
beta0 = 20.0
A1 = np.asarray(jax.jacfwd(lambda u: conv_diff(beta0, u))(jnp.zeros(n1)))   # small problem: an explicit matrix as the reference
lam_ref = np.linalg.solve(A1.T, np.asarray(w1))                             # dJ/df = A^{-T} w
assert not np.allclose(A1, A1.T)

# ❌ Wrong: declaring symmetric=True for a nonsymmetric operator; the reverse-mode gradient is silently wrong. Leaving out transpose_solve raises an error outright
def J_sym_wrong(f):
    u = lax.custom_linear_solve(lambda v: conv_diff(beta0, v), f, gmres_solve, symmetric=True)
    return jnp.vdot(w1, u)

def J_no_transpose(f):
    u = lax.custom_linear_solve(lambda v: conv_diff(beta0, v), f, gmres_solve)
    return jnp.vdot(w1, u)

try:
    err_sym = rel(jax.grad(J_sym_wrong)(f1), lam_ref)
    print(f"symmetric=True: relative gradient error {err_sym:.3f} (no warning at all)")
except Exception as e:
    err_sym = None
    print(type(e).__name__, e)
try:
    jax.grad(J_no_transpose)(f1)
    msg_tr = None
except TypeError as e:
    msg_tr = str(e)
    print("Without transpose_solve: TypeError:", msg_tr)
assert err_sym > 0.1
assert msg_tr is not None and "transpose_solve" in msg_tr

# ✅ Correct: symmetric=False and provide transpose_solve (here also using GMRES to solve A^T λ = b)
def u_cd(beta, f):
    return lax.custom_linear_solve(lambda v: conv_diff(beta, v), f, gmres_solve,
                                   transpose_solve=gmres_solve, symmetric=False)

J_cd = lambda beta, f: jnp.vdot(w1, u_cd(beta, f))
g_f = jax.grad(J_cd, argnums=1)(beta0, f1)
g_beta = jax.grad(J_cd, argnums=0)(beta0, f1)
fd_beta = (J_cd(beta0 + 1e-5, f1) - J_cd(beta0 - 1e-5, f1)) / 2e-5
print(f"dJ/df relative error {rel(g_f, lam_ref):.1e}; dJ/dβ: AD {float(g_beta):.10f}  finite difference {float(fd_beta):.10f}")
assert rel(g_f, lam_ref) < 1e-10
assert abs(g_beta - fd_beta) < 1e-8 * abs(g_beta)
# Forward mode only uses solve, so it isn't affected by the symmetric declaration
assert abs(jax.jvp(lambda b: J_cd(b, f1), (beta0,), (1.0,))[1] - g_beta) < 1e-10

src = inspect.getsource(sys.modules[jsl.cg.__module__]._isolve)
assert "custom_linear_solve" in src and "transpose_solve" in src
print("The internal implementation of jsl.cg calls custom_linear_solve ✓")

def J_jsl(k, f, u_obs, h):
    u = jsl.cg(lambda v: apply_A(k, v, h), f, tol=1e-12, maxiter=10 * f.size)[0]
    return 0.5 * h**2 * jnp.sum((u - u_obs) ** 2)

g_jsl = jax.jit(jax.grad(J_jsl))(k0, f_src, u_obs, h)
print(f"jsl.cg gradient vs hand-written custom_linear_solve: relative difference {rel(g_jsl, g_imp):.1e}; "
      f"temporary memory {temp_mib(J_jsl, k0, f_src, u_obs, h):.3f} MiB")
assert rel(g_jsl, g_imp) < 1e-9
assert temp_mib(J_jsl, k0, f_src, u_obs, h) < MEM["unrolled 100 steps"] / 10

def F_nl(u, k, c, f, h):
    return apply_A(k, u, h) + c * u**3 - f

def newton_krylov(F_u, u0, n_newton=12):
    """Newton's method, with Jacobian-vector products from jvp and inner CG (the Jacobian is symmetric positive definite)"""
    def body(i, u):
        r = F_u(u)
        Jv = lambda v: jax.jvp(F_u, (u,), (v,))[1]
        return u + cg_while(Jv, -r, tol=1e-13)
    return lax.fori_loop(0, n_newton, body, u0)

f_nl = 20.0 * jnp.ones((N, N))

def L_root(k, c, tangent_solve):
    u = lax.custom_root(lambda u: F_nl(u, k, c, f_nl, h), jnp.zeros((N, N)), newton_krylov, tangent_solve)
    return h**2 * jnp.sum(u**2)

u_nl = lax.custom_root(lambda u: F_nl(u, k0, 5.0, f_nl, h), jnp.zeros((N, N)), newton_krylov,
                       lambda g, y: lax.custom_linear_solve(g, y, cg_while, symmetric=True))
res_nl = float(jnp.linalg.norm(F_nl(u_nl, k0, 5.0, f_nl, h)) / jnp.linalg.norm(f_nl))
print(f"Newton–Krylov converged: relative residual {res_nl:.1e}")
assert res_nl < 1e-11

# ❌ Wrong: calling a while_loop solver directly inside tangent_solve (not transposable), or calling jsl.cg directly (errors on 0.12.0.1)
bad_tangent = {
    "plain cg_while": lambda g, y: cg_while(g, y),
    "plain jsl.cg":   lambda g, y: jsl.cg(g, y, tol=1e-13)[0],
}
bad_results = {}
for name, ts in bad_tangent.items():
    try:
        jax.grad(L_root, argnums=1)(k0, 5.0, ts)
        bad_results[name] = "no error"
    except Exception as e:
        bad_results[name] = type(e).__name__
        print(f"{name:<14} → {type(e).__name__}: {str(e).splitlines()[0][:90]}")
    # Forward mode doesn't need to transpose tangent_solve, so both work there
    _, t = jax.jvp(lambda c: L_root(k0, c, ts), (5.0,), (1.0,))
    print(f"{'':<14}   but jvp works: dL/dc = {float(t):.8f}")
assert bad_results["plain cg_while"] == "ValueError"
assert bad_results["plain jsl.cg"] == "NotImplementedError"