| """Preprocess MPM trajectories (L1/L2) into chart states for training: |
| per frame: Kabsch rigid pose (t, R) vs the torus rest cloud -> body-frame cloud -> CHART INVERSION |
| (warm-started Adam over the KQ=12 chart coords) -> xi(t). Stores per-trajectory npz with |
| (t, R, xi, inversion residual per frame). The MPM SUB particle indices index the SAME 381,414-particle |
| cloud the chart was trained on, so correspondence is exact. |
| Note (honest): fps=24 makes FD accelerations at impact frames coarse — the A' residual masks |
| high-jerk frames; short-window recon (loss C) carries contact. v2 data can re-export at fps=120. |
| Usage: prep_mpm_traj.py L1|L2 [GPU-visible via env]""" |
| import glob |
| import os |
| import sys |
| import time |
| import numpy as np |
| import torch |
|
|
| sys.path.insert(0, "/home/qiuyid/neurok/experiments/repro") |
| from neurok_faithful import NeurokFaithful, K as KTOK, DMODEL, make_functorch_safe, dq_apply_t |
|
|
| DEV = "cuda"; OUT = "/scr/qiuyid/dingqy/out"; DATA = os.environ.get("MPMDATA", "/home/qiuyid/neurok/data/mpm_torus") |
| MODE = sys.argv[1] if len(sys.argv) > 1 else "L1" |
| KQ = 12; NINV = 4096; INV_STEPS = 120; INV_STEPS_WARM = int(os.environ.get("WARM", "40")) |
| torch.manual_seed(0) |
|
|
| m = NeurokFaithful().to(DEV) |
| m.load_state_dict(torch.load(os.environ.get("CHARTCK", f"{OUT}/neurok_torus.pt"), map_location=DEV)); m.eval() |
| make_functorch_safe(m.dec) |
| for p in m.parameters(): |
| p.requires_grad_(False) |
| CK = torch.load(os.environ.get("CODESCK", f"{OUT}/torus_codes.pt"), weights_only=False) |
| zf = CK["z"].to(DEV); rest = torch.tensor(CK["rest"], device=DEV) |
| z_rest = zf[0].reshape(-1) |
| Adev = (zf[1:].reshape(len(zf) - 1, -1) - z_rest[None]) |
| _, S_, V_ = torch.linalg.svd(Adev, full_matrices=False) |
| Q = V_[:KQ].contiguous() |
| span = float((rest.max(0).values - rest.min(0).values).mean()) |
|
|
|
|
| def chart_c(xi, P): |
| y = dq_apply_t(m.dec((z_rest + xi @ Q).view(1, KTOK, DMODEL), P[None])[0], P) |
| return y - y.mean(0, keepdim=True) |
|
|
|
|
| def kabsch(A, B): |
| H = (A - A.mean(0)).T @ (B - B.mean(0)); U, S, Vt = np.linalg.svd(H) |
| return (U @ np.diag([1, 1, np.sign(np.linalg.det(U @ Vt))]) @ Vt).astype(np.float32) |
|
|
|
|
| files = sorted(glob.glob(f"{DATA}/{MODE}_shard*_traj*.npz")) |
| print(f"{MODE}: {len(files)} trajectories to preprocess", flush=True) |
| gsel = np.random.RandomState(3).choice(30000, NINV, replace=False) |
|
|
| t0 = time.time() |
| for fi, fpath in enumerate(files): |
| opath = fpath.replace(".npz", "_prep.npz") |
| if os.path.exists(opath) or fpath.endswith("_prep.npz"): |
| continue |
| d = np.load(fpath) |
| traj = d["traj"].astype(np.float32) |
| sub_idx = d["sub_idx"] |
| Psub = rest[torch.tensor(sub_idx[gsel], device=DEV)].contiguous() |
| rest_sub = rest[torch.tensor(sub_idx, device=DEV)].cpu().numpy() |
| F = traj.shape[0] |
| ts, Rs, xis, errs = [], [], [], [] |
| xi = torch.zeros(KQ, device=DEV, requires_grad=True) |
| for f in range(F): |
| Rk = kabsch(rest_sub, traj[f]).T |
| tk = traj[f].mean(0) |
| yk = (traj[f] - tk) @ Rk |
| ytgt = torch.tensor(yk[gsel], device=DEV) |
| ytgt = ytgt - ytgt.mean(0, keepdim=True) |
| opt = torch.optim.Adam([xi], lr=0.4 if f == 0 else 0.15) |
| steps = INV_STEPS if f == 0 else INV_STEPS_WARM |
| for it in range(steps): |
| loss = ((chart_c(xi, Psub) - ytgt) ** 2).mean() |
| opt.zero_grad(); loss.backward(); opt.step() |
| ts.append(tk); Rs.append(Rk); xis.append(xi.detach().cpu().numpy()); errs.append(float(loss.detach().sqrt()) / span) |
| np.savez_compressed(opath + ".tmp.npz", t=np.stack(ts), R=np.stack(Rs), xi=np.stack(xis), |
| inv_err=np.array(errs), E=d["E"], nu=d["nu"], vel=d["vel"], omega=d["omega"], fps=d["fps"]) |
| os.replace(opath + ".tmp.npz", opath) |
| if fi % 10 == 0: |
| el = time.time() - t0 |
| print(f" {fi+1}/{len(files)}: inv_err mean {np.mean(errs)*100:.2f}% span, {el/60:.1f} min, " |
| f"ETA {(el/max(fi+1,1))*(len(files)-fi-1)/60:.0f} min", flush=True) |
| print("PREP DONE", flush=True) |
|
|
