import json, math, random, time from pathlib import Path import numpy as np import torch import torch.nn as nn from torch.func import jacrev, vmap SEED = 1234 random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED) torch.set_default_dtype(torch.float32) device = "cuda" if torch.cuda.is_available() else "cpu" def cofactor(a): # Stable 3x3 cofactor via cross products; unlike det(A) A^{-T}, # this remains finite near singular matrices. c0 = torch.cross(a[..., :, 1], a[..., :, 2], dim=-1) c1 = torch.cross(a[..., :, 2], a[..., :, 0], dim=-1) c2 = torch.cross(a[..., :, 0], a[..., :, 1], dim=-1) return torch.stack((c0, c1, c2), dim=-1) def polar_so3(a): u, _, vh = torch.linalg.svd(a, full_matrices=False) q = u @ vh sign = torch.where(torch.linalg.det(q) < 0, -torch.ones_like(q[..., 0, 0]), torch.ones_like(q[..., 0, 0])) # Flip the last right-singular vector when the polar factor is a reflection. vh = vh.clone() vh[..., -1, :] = vh[..., -1, :] * sign[..., None] return u @ vh def target_map(x, amp=1.0): # Smooth orientation-preserving twist with radial compression. z = x[..., 2] th = amp * 1.8 * z c, s = torch.cos(th), torch.sin(th) xx, yy = 0.68 * x[..., 0], 0.68 * x[..., 1] return torch.stack((c*xx - s*yy, s*xx + c*yy, x[..., 2]), -1) class Warp(nn.Module): def __init__(self, with_rotation=False): super().__init__() self.with_rotation = with_rotation self.body = nn.Sequential(nn.Linear(3, 48), nn.Tanh(), nn.Linear(48, 48), nn.Tanh()) self.phi = nn.Linear(48, 3) self.rot = nn.Linear(48, 3) if with_rotation else None # Start close to identity, but retain trainability. nn.init.zeros_(self.phi.weight); nn.init.zeros_(self.phi.bias) if self.rot is not None: nn.init.zeros_(self.rot.weight); nn.init.zeros_(self.rot.bias) def forward(self, x): h = self.body(x) y = self.phi(h) + x if self.with_rotation: # Stable SO(3) frame: exponential map avoids SVD gradients at # repeated singular values of the identity initialization. w = self.rot(h) z = torch.zeros_like(w[..., 0]) wx = torch.stack((z, -w[..., 2], w[..., 1], w[..., 2], z, -w[..., 0], -w[..., 1], w[..., 0], z), dim=-1) wx = wx.reshape(*w.shape[:-1], 3, 3) return y, torch.matrix_exp(wx) return y def jacobian(model, x): # Per-point forward deformation Jacobian; create_graph enables regularizer training. if model.with_rotation: f = lambda z: model(z)[0] else: f = model return vmap(jacrev(f))(x) def identity_check(): g = torch.Generator(device=device).manual_seed(SEED + 10) a = torch.randn(8, 3, 3, device=device, generator=g) q, _ = torch.linalg.qr(torch.randn(8, 3, 3, device=device, generator=g)) # QR can contain reflections; make proper rotations. q[..., :, -1] *= torch.where(torch.linalg.det(q) < 0, -1., 1.)[:, None] lhs_c = cofactor(q.transpose(-1, -2) @ a) rhs_c = q.transpose(-1, -2) @ cofactor(a) lhs_d = torch.linalg.det(q.transpose(-1, -2) @ a) rhs_d = torch.linalg.det(a) return {"cofactor_max_abs": float((lhs_c-rhs_c).abs().max()), "det_max_abs": float((lhs_d-rhs_d).abs().max()), "rotation_det_max_abs": float((torch.linalg.det(q)-1).abs().max())} def run_variant(kind, amp, steps=500): torch.manual_seed(SEED + int(amp*10) + (0 if kind == "baseline" else 100)) model = Warp(with_rotation=(kind == "polyconvex")).to(device) opt = torch.optim.Adam(model.parameters(), lr=1e-3) # Fixed collocation set gives an equal-step comparison. gen = torch.Generator(device=device).manual_seed(SEED + 77) x = (torch.rand(96, 3, device=device, generator=gen) * 2 - 1).requires_grad_(True) y = target_map(x, amp).detach() lam = 0.10 det_lam = 0.015 start = time.time(); history = [] for step in range(steps): opt.zero_grad(set_to_none=True) if kind == "baseline": pred = model(x); J = jacobian(model, x) data = ((pred-y)**2).mean() reg = ((J-torch.eye(3, device=device))**2).mean() d = torch.linalg.det(J) # Ordinary baseline includes the same basic positive-Jacobian barrier. barrier = (-torch.log(torch.clamp(d, min=1e-4)) + 20.0*torch.relu(-d)**2).mean() loss = data + lam*reg + det_lam*barrier else: pred, R = model(x); J = jacobian(model, x) data = ((pred-y)**2).mean() U = R.transpose(-1,-2) @ J C = cofactor(U); d = torch.linalg.det(J) # Explicit convex quadratic in lifted variables plus determinant barrier. eye = torch.eye(3, device=device) pc = ((U-eye)**2).mean() + 0.5*((C-eye)**2).mean() + 0.5*((d-1)**2).mean() barrier = (-torch.log(torch.clamp(d, min=1e-4)) + 20.0*torch.relu(-d)**2).mean() loss = data + lam*pc + det_lam*barrier if not torch.isfinite(loss): break loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step() if step in (0, 99, 249, steps-1): history.append(float(data.detach())) if kind == "polyconvex": pred, R = model(x) else: pred = model(x) J = jacobian(model, x).detach(); d = torch.linalg.det(J) err = ((pred-y)**2).mean().item() fold = (d <= 0).float().mean().item() cond = torch.linalg.cond(J).nan_to_num(1e6).mean().item() dmin = d.min().item() return {"final_task_mse": err, "fold_fraction": fold, "mean_condition": cond, "min_det": dmin, "task_history": history, "seconds": time.time()-start} def main(): global device checks = identity_check() results = {"device": device, "identity_check": checks, "runs": {}} for amp in (1.0, 2.0): for kind in ("baseline", "polyconvex"): try: results["runs"][f"{kind}_amp{amp}"] = run_variant(kind, amp) except Exception as e: # CUDA OOM or unsupported autodiff falls back to a CPU rerun. if device == "cuda": device = "cpu" torch.set_default_device("cpu") results["device_fallback"] = str(e) results["runs"][f"{kind}_amp{amp}"] = run_variant(kind, amp) else: raise Path("results.json").write_text(json.dumps(results, indent=2)) print(json.dumps(results, indent=2)) if __name__ == "__main__": main()