masiv/scripts/verify_plas.py

"""Charter-mandated verification for the plasticity module (run BEFORE any training uses PLAS=1).
T1  SafeSVD3 backward vs finite differences — random matrices AND the dangerous repeated-singular-
    value cases (uniform squash sigma1=sigma2) where torch's own SVD backward blows up.
T2  return_map analytic: uniaxial stretch beyond y_hi must clamp exactly and produce the analytic
    plastic split Fp = diag(a/y_hi, 1, 1); in-window states must pass through BIT-IDENTICALLY and
    keep Fp = I (load-unload => permanent set retained in Fp).
T3  (documented, run at egg acceptance) GRID A/B: PLAS=1 (window open) vs PLAS=0 FIXQ=1 must give
    identical CRN losses — the whole-pipeline no-op regression.
"""
import torch

torch.manual_seed(0)
DEV = "cuda" if torch.cuda.is_available() else "cpu"
import sys

sys.path.insert(0, "/home/qiuyid/neurok/experiments/idea1/design4")
from plastic import SafeSVD3, return_map

ok = True

# ---- T1: return-map gradient vs FD on SIGN-INVARIANT observables ----
# U/Vh columns have arbitrary signs that flip between nearby SVD calls, so FD on U/Vh directly is
# ill-posed; the objects training differentiates — the clamped reconstruction B = U clamp(S) Vh and
# the singular values S — are invariant, so FD on them is well-defined even at repeated sigmas.
YLO = torch.tensor(0.8, dtype=torch.float64, device=DEV)
YHI = torch.tensor(1.2, dtype=torch.float64, device=DEV)


def recon_scalar(A, w):
    B, _ = return_map(A, YLO, YHI)
    return (B * w).sum()


def sv_scalar(A, wS):
    _, S, _ = SafeSVD3.apply(A)
    return (S * wS).sum()


cases = {
    "random": torch.randn(8, 3, 3, dtype=torch.float64, device=DEV) * 0.5 + torch.eye(3, dtype=torch.float64, device=DEV),
    "big-squash(clamps)": torch.stack([torch.diag(torch.tensor(d, dtype=torch.float64, device=DEV))
                                       for d in ([0.5, 0.5, 1.6], [0.4, 0.9, 0.9], [1.5, 1.5, 0.5],
                                                 [0.5, 0.5001, 1.6], [0.3, 1.0, 1.0], [1.4, 1.4, 1.4])]),
}
for name, A0 in cases.items():
    A0 = A0 + 1e-2 * torch.randn_like(A0)                                     # generic rotations, keeps near-repeated sigmas
    w = torch.randn(A0.shape[0], 3, 3, dtype=torch.float64, device=DEV)
    wS = torch.randn(A0.shape[0], 3, dtype=torch.float64, device=DEV)
    for tag, fn, ww in (("recon", recon_scalar, w), ("svals", sv_scalar, wS)):
        A = A0.clone().requires_grad_(True)
        g = torch.autograd.grad(fn(A, ww), A)[0]
        eps = 1e-6
        errs = []
        for b, i, j in [(0, 0, 0), (1, 1, 2), (2, 2, 1), (3, 0, 1), (4, 2, 0), (5, 1, 1)]:
            Ap = A0.clone(); Ap[b, i, j] += eps
            Am = A0.clone(); Am[b, i, j] -= eps
            fd = float(fn(Ap, ww) - fn(Am, ww)) / (2 * eps)
            errs.append(abs(fd - float(g[b, i, j])) / (abs(fd) + 1e-6))
        worst = max(errs)
        print(f"T1 {tag}-vs-FD [{name}]: worst rel err {worst:.2e} finite={bool(torch.isfinite(g).all())}", flush=True)
        ok &= worst < 1e-3 and bool(torch.isfinite(g).all())

# ---- T2: analytic uniaxial box clamp ----
ylo = torch.tensor(0.8, dtype=torch.float64, device=DEV, requires_grad=True)
yhi = torch.tensor(1.2, dtype=torch.float64, device=DEV, requires_grad=True)
a = 1.5
F = torch.diag(torch.tensor([a, 1.0, 1.0], dtype=torch.float64, device=DEV))[None]
Fe, need = return_map(F, ylo, yhi)
Fp = torch.linalg.inv(Fe) @ F
an_Fe = torch.diag(torch.tensor([1.2, 1.0, 1.0], dtype=torch.float64, device=DEV))
an_Fp = torch.diag(torch.tensor([a / 1.2, 1.0, 1.0], dtype=torch.float64, device=DEV))
e1 = float((Fe[0] - an_Fe).abs().max()); e2 = float((Fp[0] - an_Fp).abs().max())
print(f"T2 uniaxial clamp: |Fe-analytic| {e1:.2e}  |Fp-analytic| {e2:.2e}  clamped={bool(need[0])}", flush=True)
ok &= e1 < 1e-9 and e2 < 1e-9 and bool(need[0])
gy = torch.autograd.grad(Fe.sum(), yhi, retain_graph=False)[0]                # d(clamped sigma)/d yhi = 1 exactly
print(f"T2 dFe/dyhi = {float(gy):.6f} (analytic 1.0)", flush=True)
ok &= abs(float(gy) - 1.0) < 1e-6
# in-window passthrough: bit-identical object, Fp stays I after unload
F2 = torch.diag(torch.tensor([1.1, 0.95, 1.0], dtype=torch.float64, device=DEV))[None]
Fe2, need2 = return_map(F2, ylo, yhi)
print(f"T2 in-window passthrough: same-object={Fe2 is F2} clamped={bool(need2[0])}", flush=True)
ok &= (Fe2 is F2) and not bool(need2[0])

print("VERIFY_PLAS:", "ALL PASS" if ok else "FAIL", flush=True)