第 11 章　科学计算与可微分仿真

# 第一件事：开启 64 位精度（必须在创建任何 JAX 数组之前）
import jax
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp

# 自证：默认浮点类型、随机数、Python 标量转换都是 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("默认浮点 dtype:", jnp.zeros(3).dtype)

# 公共导入
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("设备:", jax.devices())

def rel(a, b):
    """相对误差 ||a-b|| / ||b||（全部转成 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):
    """先预热（编译），再计时 n 次取平均，单位毫秒。只用于打印，不做 assert。"""
    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

# ❌ 错误写法：先创建数组、后开启 x64 —— 已有数组仍是 float32，且会经由弱类型运算传染
probe = textwrap.dedent("""
    import jax, jax.numpy as jnp, json
    a = jnp.ones(3)                                  # 此时默认还是 float32
    jax.config.update("jax_enable_x64", True)        # 太晚了
    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("实测 dtype:", late)
except Exception as e:
    late = None
    print("子进程失败:", type(e).__name__, e)
assert late["a"] == "float32" and late["a*1.0"] == "float32"   # 旧数组仍是单精度，且乘标量后还是单精度
assert late["b"] == "float64"

# ✅ 正确写法：先 update 再创建任何数组（同样在子进程里验证；本内核第 0 节就是这样做的）
probe_ok = textwrap.dedent("""
    import jax
    jax.config.update("jax_enable_x64", True)        # 第一件事
    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("实测 dtype:", 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   # 本内核

# ❌ 错误写法：float32 下做 Poisson 收敛性研究 —— 细网格上舍入误差淹没离散误差，二阶收敛「消失」
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))           # 随机右端项（不是特征向量）
    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 最大误差:", np.array2string(e32, precision=3))
    print("相邻网格误差比:", np.round(e32[:-1] / e32[1:], 3), "  （二阶收敛应约等于 4）")
    print(f"float32 有限差分梯度 {res32['fd']:.6f}  vs  AD 梯度 {res32['ad']:.6f}"
          f"  相对误差 {abs(res32['fd'] - res32['ad']) / abs(res32['ad']):.1e}")
    print(f"float32 有限差分（步长 1e-3，接近最优）相对误差 {abs(res32['fd_best'] - res32['ad']) / abs(res32['ad']):.1e}"
          f"；CG（tol=1e-10，N=63 随机右端项）的真实相对残差 {res32['res63']:.1e}")
except Exception as e:
    res32 = None
    print("子进程失败:", type(e).__name__, e)
assert e32[-2] / e32[-1] < 3.5                                     # 最细一级的误差比明显偏离 4
assert abs(res32["fd"] - res32["ad"]) / abs(res32["ad"]) > 1e-3    # 步长 1e-6 时 float32 有限差分只剩 2～3 位有效数字
assert 1e-6 < abs(res32["fd_best"] - res32["ad"]) / abs(res32["ad"]) < 1e-4   # 最优步长附近也只有 4～5 位
assert res32["res63"] > 1e-7                                       # CG 的真实残差降不到 1e-7

# ✅ 正确写法：float64 下做梯度检验（收敛性研究见第 2 节，二阶收敛一直保持到最细网格）
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 有限差分 {float(fd64):.10f}  vs  AD {float(ad64):.10f}  相对误差 {abs(fd64 - ad64) / abs(ad64):.1e}")
assert abs(fd64 - ad64) / abs(ad64) < 1e-8
# float64 下的 Poisson 收敛阶见第 2 节

def neg_lap(u, h):
    """无矩阵五点差分算子 -Δ_h u（齐次 Dirichlet 边界，用 pad 补 0 实现）"""
    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 参考：用 Kronecker 积组装同一个矩阵（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)            # 行优先展平，与 kron(I, T) + kron(T, I) 的排列一致
err_pad = rel(neg_lap(jnp.asarray(v_np), h), ref)
print(f"pad 切片算子 vs SciPy CSR：相对误差 {err_pad:.1e}")
assert err_pad < 1e-14

# ❌ 错误写法：直接用 jnp.roll 取邻点 —— roll 会把对边的值卷回来，等于周期边界条件
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"直接 roll 的算子 vs SciPy（Dirichlet）：整体相对误差 {err_roll:.3f}；"
          f"内部节点 {err_inner:.1e}，第一行 {err_edge:.3f} —— 不报错，但贴边的节点全错")
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

# ✅ 正确写法：先 pad 一圈 0，再在扩充后的数组上 roll，最后裁掉外圈（与切片写法完全等价）
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 的算子 vs SciPy：相对误差 {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，齐次 Dirichlet；算子以函数形式传给 cg"""
    u, info = jsl.cg(lambda v: neg_lap(v, h), f, tol=1e-12, maxiter=10 * f.size)
    return u

def u_star(p):
    """制造解 u*(x, y)，在单位正方形边界上为 0"""
    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*，由 AD 自动求出

# 抽查一个点：与手推的解析式对比
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):
    """返回 h、网格坐标、制造解 u* 和源项 f = -Δu*（都在 N×N 内部节点上）"""
    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} {'最大误差':>12} {'相对残差':>10} {'耗时(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                                # 代数误差远小于离散误差

errs = np.array(errs)
ratios = errs[:-1] / errs[1:]
orders = np.log2(ratios)
print("误差比:", np.round(ratios, 4))
print("收敛阶:", np.round(orders, 4))
assert np.all(np.abs(ratios - 4) < 0.1)                   # 网格加密一倍，误差缩小到约 1/4
assert np.all(np.abs(orders - 2) < 0.02)                  # 二阶收敛
POISSON_ORDERS = orders

# 与 SciPy 稀疏直接求解对拍（同一个离散系统，应在舍入误差内一致）
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：相对差 {rel(u_jax, u_sp):.1e}")
assert rel(u_jax, u_sp) < 1e-10

# 验证：info 是 None；maxiter 太小时静默返回未收敛的解
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, "| 只迭代 5 步时的相对残差:", 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 三元组 (row, col, val)，NumPy 生成（前处理不需要可微）"""
    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 参考
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 作为参数传入，不在闭包里
y_bcoo = matvec(A_bcoo, jnp.asarray(x_np))
y_csr = A_csr @ x_np
print(f"BCOO matvec vs SciPy CSR：相对误差 {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   # 与无矩阵算子一致

# 从 SciPy 直接转换，以及转回稠密做小规模检查
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())

# 有限元式组装：重复索引会累加（BCOO 保留重复项，matvec 时相加；sum_duplicates 合并）
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("稠密形式:\n", A_dup.todense(), "\nnse:", A_dup.nse, "→ 合并后 nse:", 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):
    """照着稠密矩阵写：||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 结果 {float(e_sp):.6e}  vs  SciPy {e_ref:.6e}")
assert abs(float(e_sp) - e_ref) / abs(e_ref) < 1e-13

# 对矩阵非零元求导：d/dA_ij ||Ax||^2 = 2 (Ax)_i x_j，只在非零位置上取值
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("对 data 的梯度形状:", g_data.shape, "| 与解析式相对误差:", 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：相对差 {rel(u_cg_bcoo, u_ref):.1e}")
assert rel(u_cg_bcoo, u_ref) < 1e-10

A_bcsr = sparse.BCSR.from_bcoo(A_bcoo)                  # 转成 CSR 三件套
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 回调 SciPy）vs SciPy：{rel(u_ss, u_ref):.1e}；可以 jit、可以 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   # 梯度 = A^{-T} 1
assert "scipy.sparse.linalg.spsolve" in inspect.getsource(sparse.linalg._spsolve_cpu_lowering)

# ❌ 错误写法：对 spsolve 做 vmap（一次解多个右端项）—— 没有 batching 规则
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

# ✅ 正确写法：批量右端项用 vmap(cg)（matvec 是 BCOO，可批处理）；或者在 SciPy 里一次分解、多次回代
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：一次 LU，多次回代
U_ref = np.stack([lu.solve(np.asarray(bb)) for bb in B])
print(f"vmap(cg) 批量求解 vs SciPy splu：{rel(U, U_ref):.1e}")
assert rel(U, U_ref) < 1e-10

# 计时（只打印）：同一个 N=63 的五点算子，三种 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"无矩阵 matvec  : {t_mf * 1e3:8.1f} µs")
print(f"SciPy CSR      : {t_sp * 1e3:8.1f} µs   (NumPy 侧，无 JAX 调度开销)")

def apply_A(k, u, h):
    """无矩阵变系数扩散算子 -∇·(k∇u)，k 和 u 都是 N×N 节点数组，边界外的 k 取最近值"""
    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])            # x 方向面系数 (N+1, N)
    ky = 0.5 * (kp[1:-1, 1:] + kp[1:-1, :-1])            # y 方向面系数 (N, N+1)
    fx = kx * (up[1:, 1:-1] - up[:-1, 1:-1])              # 面通量
    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              # 正定性抽查
print("变系数算子：退化检查、对称性、linear_transpose 一致性全部通过")

def cg_unrolled(matvec, b, n_iter):
    """固定 n_iter 步的 CG（静态步数 fori_loop → 可反向求导，但要存每一步）"""
    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（while_loop，不能直接反向求导）"""
    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):
    """隐式微分：AD 只用到 matvec（对闭包里的 k 求导）和 solve 本身"""
    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)       # 一个人造的「观测」

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)

# 有限差分（中心差分，沿 3 个随机方向），前向用紧容差的 cg_while
J_fd = jax.jit(J_implicit)
eps = 1e-5
print(f"{'方向':>4} {'有限差分':>16} {'隐式微分':>16} {'展开求导':>16} {'隐式相对误差':>12} {'展开相对误差':>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                                       # 隐式梯度与有限差分一致
    assert eu < 1e-3                                       # 展开 300 步的梯度也与有限差分一致（但差了好几个数量级）

GRAD_REL_UNR_IMP = rel(g_unr, g_imp)
print(f"\n展开 {K_UNROLL} 步 vs 隐式：梯度相对差 {GRAD_REL_UNR_IMP:.2e}")
assert GRAD_REL_UNR_IMP < 1e-4

# 内存：读取 XLA 编译后反向传播所需的临时缓冲区
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"展开 {K} 步"] = temp_mib(J_unrolled, k0, f_src, u_obs, h, K, static_argnums=4)
MEM["隐式 (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["展开 300 步"] > 2.5 * MEM["展开 100 步"]          # 随迭代次数线性增长
assert MEM["隐式 (custom_linear_solve)"] < MEM["展开 300 步"] / 50

print("\n计时（只打印）：")
print(f"  展开 300 步 grad: {timeit(grad_unr, k0, f_src, u_obs, h, K_UNROLL):.2f} ms")
print(f"  隐式 grad       : {timeit(grad_imp, k0, f_src, u_obs, h):.2f} ms")

print("\n隐式版本保存的 residuals（只有输入 k 和由解 u 算出的几个同规模数组，没有任何按迭代步堆叠的量）：")
print_saved_residuals(lambda k: J_implicit(k, f_src, u_obs, h), k0)

print(f"{'K':>5} {'前向解误差':>12} {'展开梯度误差':>14} {'隐式梯度误差':>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：恰好迭代 K 步
    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)

# 同一个 K=80：展开函数自身的有限差分 = 真梯度；但 AD 的 jvp（前向）和 vjp（反向）都远离真值，彼此也不一致
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：真梯度 {float(true80):.6e}｜展开函数自身的有限差分 {float(fd80):.6e}｜"
      f"展开 jvp {float(jvp80):.6e}｜展开 vjp {float(vjp80):.6e}")
assert abs(fd80 - true80) < 1e-6 * abs(fd80)               # 函数值本身已收敛，其有限差分就是真梯度
assert abs(jvp80 - true80) > 100 * abs(true80)             # 前向模式的展开导数同样失真 → 不是反向模式的 bug
assert abs(vjp80 - true80) > 100 * abs(true80)

assert scan_rows[:, 2].max() > 1e-2                          # 展开梯度在收敛阶段会严重失真
assert np.all(scan_rows[scan_rows[:, 0] >= 80][:, 3] < 1e-5)  # 隐式梯度随前向一起稳定收敛

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

def conv_diff(beta, u):
    """一维对流扩散 -u'' + beta u'（中心差分，Dirichlet 0），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)))   # 小规模：显式矩阵做参考
lam_ref = np.linalg.solve(A1.T, np.asarray(w1))                             # dJ/df = A^{-T} w
assert not np.allclose(A1, A1.T)

# ❌ 错误写法：非对称算子却声明 symmetric=True —— 反向梯度静默算错；不给 transpose_solve 则直接报错
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：梯度相对误差 {err_sym:.3f}（无任何警告）")
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("不给 transpose_solve：TypeError:", msg_tr)
assert err_sym > 0.1
assert msg_tr is not None and "transpose_solve" in msg_tr

# ✅ 正确写法：symmetric=False 并提供 transpose_solve（这里同样用 GMRES 解 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 相对误差 {rel(g_f, lam_ref):.1e}；dJ/dβ: AD {float(g_beta):.10f}  有限差分 {float(fd_beta):.10f}")
assert rel(g_f, lam_ref) < 1e-10
assert abs(g_beta - fd_beta) < 1e-8 * abs(g_beta)
# 前向模式只用 solve，不受 symmetric 声明影响
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("jsl.cg 的内部实现调用了 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 的梯度 vs 手写 custom_linear_solve：相对差 {rel(g_jsl, g_imp):.1e}；"
      f"临时内存 {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["展开 100 步"] / 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 法，Jacobian-向量积用 jvp，内层 CG（Jacobian 对称正定）"""
    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 收敛：相对残差 {res_nl:.1e}")
assert res_nl < 1e-11

# ❌ 错误写法：tangent_solve 里直接调用 while_loop 求解器（不可转置），或直接调用 jsl.cg（0.12.0.1 上会报错）
bad_tangent = {
    "直接用 cg_while": lambda g, y: cg_while(g, y),
    "直接用 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] = "无报错"
    except Exception as e:
        bad_results[name] = type(e).__name__
        print(f"{name:<14} → {type(e).__name__}: {str(e).splitlines()[0][:90]}")
    # 前向模式不需要转置 tangent_solve，两者都能算
    _, t = jax.jvp(lambda c: L_root(k0, c, ts), (5.0,), (1.0,))
    print(f"{'':<14}   但 jvp 正常：dL/dc = {float(t):.8f}")
assert bad_results["直接用 cg_while"] == "ValueError"
assert bad_results["直接用 jsl.cg"] == "NotImplementedError"