1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6from torch.func import jacrev, vmap
  7
  8SEED = 1234
  9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 10torch.set_default_dtype(torch.float32)
 11device = "cuda" if torch.cuda.is_available() else "cpu"
 12
 13
 14def cofactor(a):
 15    # Stable 3x3 cofactor via cross products; unlike det(A) A^{-T},
 16    # this remains finite near singular matrices.
 17    c0 = torch.cross(a[..., :, 1], a[..., :, 2], dim=-1)
 18    c1 = torch.cross(a[..., :, 2], a[..., :, 0], dim=-1)
 19    c2 = torch.cross(a[..., :, 0], a[..., :, 1], dim=-1)
 20    return torch.stack((c0, c1, c2), dim=-1)
 21
 22
 23def polar_so3(a):
 24    u, _, vh = torch.linalg.svd(a, full_matrices=False)
 25    q = u @ vh
 26    sign = torch.where(torch.linalg.det(q) < 0, -torch.ones_like(q[..., 0, 0]), torch.ones_like(q[..., 0, 0]))
 27    # Flip the last right-singular vector when the polar factor is a reflection.
 28    vh = vh.clone()
 29    vh[..., -1, :] = vh[..., -1, :] * sign[..., None]
 30    return u @ vh
 31
 32
 33def target_map(x, amp=1.0):
 34    # Smooth orientation-preserving twist with radial compression.
 35    z = x[..., 2]
 36    th = amp * 1.8 * z
 37    c, s = torch.cos(th), torch.sin(th)
 38    xx, yy = 0.68 * x[..., 0], 0.68 * x[..., 1]
 39    return torch.stack((c*xx - s*yy, s*xx + c*yy, x[..., 2]), -1)
 40
 41
 42class Warp(nn.Module):
 43    def __init__(self, with_rotation=False):
 44        super().__init__()
 45        self.with_rotation = with_rotation
 46        self.body = nn.Sequential(nn.Linear(3, 48), nn.Tanh(), nn.Linear(48, 48), nn.Tanh())
 47        self.phi = nn.Linear(48, 3)
 48        self.rot = nn.Linear(48, 3) if with_rotation else None
 49        # Start close to identity, but retain trainability.
 50        nn.init.zeros_(self.phi.weight); nn.init.zeros_(self.phi.bias)
 51        if self.rot is not None:
 52            nn.init.zeros_(self.rot.weight); nn.init.zeros_(self.rot.bias)
 53
 54    def forward(self, x):
 55        h = self.body(x)
 56        y = self.phi(h) + x
 57        if self.with_rotation:
 58            # Stable SO(3) frame: exponential map avoids SVD gradients at
 59            # repeated singular values of the identity initialization.
 60            w = self.rot(h)
 61            z = torch.zeros_like(w[..., 0])
 62            wx = torch.stack((z, -w[..., 2], w[..., 1],
 63                              w[..., 2], z, -w[..., 0],
 64                              -w[..., 1], w[..., 0], z), dim=-1)
 65            wx = wx.reshape(*w.shape[:-1], 3, 3)
 66            return y, torch.matrix_exp(wx)
 67        return y
 68
 69
 70def jacobian(model, x):
 71    # Per-point forward deformation Jacobian; create_graph enables regularizer training.
 72    if model.with_rotation:
 73        f = lambda z: model(z)[0]
 74    else:
 75        f = model
 76    return vmap(jacrev(f))(x)
 77
 78
 79def identity_check():
 80    g = torch.Generator(device=device).manual_seed(SEED + 10)
 81    a = torch.randn(8, 3, 3, device=device, generator=g)
 82    q, _ = torch.linalg.qr(torch.randn(8, 3, 3, device=device, generator=g))
 83    # QR can contain reflections; make proper rotations.
 84    q[..., :, -1] *= torch.where(torch.linalg.det(q) < 0, -1., 1.)[:, None]
 85    lhs_c = cofactor(q.transpose(-1, -2) @ a)
 86    rhs_c = q.transpose(-1, -2) @ cofactor(a)
 87    lhs_d = torch.linalg.det(q.transpose(-1, -2) @ a)
 88    rhs_d = torch.linalg.det(a)
 89    return {"cofactor_max_abs": float((lhs_c-rhs_c).abs().max()),
 90            "det_max_abs": float((lhs_d-rhs_d).abs().max()),
 91            "rotation_det_max_abs": float((torch.linalg.det(q)-1).abs().max())}
 92
 93
 94def run_variant(kind, amp, steps=500):
 95    torch.manual_seed(SEED + int(amp*10) + (0 if kind == "baseline" else 100))
 96    model = Warp(with_rotation=(kind == "polyconvex")).to(device)
 97    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
 98    # Fixed collocation set gives an equal-step comparison.
 99    gen = torch.Generator(device=device).manual_seed(SEED + 77)
100    x = (torch.rand(96, 3, device=device, generator=gen) * 2 - 1).requires_grad_(True)
101    y = target_map(x, amp).detach()
102    lam = 0.10
103    det_lam = 0.015
104    start = time.time(); history = []
105    for step in range(steps):
106        opt.zero_grad(set_to_none=True)
107        if kind == "baseline":
108            pred = model(x); J = jacobian(model, x)
109            data = ((pred-y)**2).mean()
110            reg = ((J-torch.eye(3, device=device))**2).mean()
111            d = torch.linalg.det(J)
112            # Ordinary baseline includes the same basic positive-Jacobian barrier.
113            barrier = (-torch.log(torch.clamp(d, min=1e-4)) + 20.0*torch.relu(-d)**2).mean()
114            loss = data + lam*reg + det_lam*barrier
115        else:
116            pred, R = model(x); J = jacobian(model, x)
117            data = ((pred-y)**2).mean()
118            U = R.transpose(-1,-2) @ J
119            C = cofactor(U); d = torch.linalg.det(J)
120            # Explicit convex quadratic in lifted variables plus determinant barrier.
121            eye = torch.eye(3, device=device)
122            pc = ((U-eye)**2).mean() + 0.5*((C-eye)**2).mean() + 0.5*((d-1)**2).mean()
123            barrier = (-torch.log(torch.clamp(d, min=1e-4)) + 20.0*torch.relu(-d)**2).mean()
124            loss = data + lam*pc + det_lam*barrier
125        if not torch.isfinite(loss):
126            break
127        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
128        if step in (0, 99, 249, steps-1): history.append(float(data.detach()))
129    if kind == "polyconvex": pred, R = model(x)
130    else: pred = model(x)
131    J = jacobian(model, x).detach(); d = torch.linalg.det(J)
132    err = ((pred-y)**2).mean().item()
133    fold = (d <= 0).float().mean().item()
134    cond = torch.linalg.cond(J).nan_to_num(1e6).mean().item()
135    dmin = d.min().item()
136    return {"final_task_mse": err, "fold_fraction": fold, "mean_condition": cond,
137            "min_det": dmin, "task_history": history, "seconds": time.time()-start}
138
139
140def main():
141    global device
142    checks = identity_check()
143    results = {"device": device, "identity_check": checks, "runs": {}}
144    for amp in (1.0, 2.0):
145        for kind in ("baseline", "polyconvex"):
146            try:
147                results["runs"][f"{kind}_amp{amp}"] = run_variant(kind, amp)
148            except Exception as e:
149                # CUDA OOM or unsupported autodiff falls back to a CPU rerun.
150                if device == "cuda":
151                    device = "cpu"
152                    torch.set_default_device("cpu")
153                    results["device_fallback"] = str(e)
154                    results["runs"][f"{kind}_amp{amp}"] = run_variant(kind, amp)
155                else:
156                    raise
157    Path("results.json").write_text(json.dumps(results, indent=2))
158    print(json.dumps(results, indent=2))
159
160if __name__ == "__main__": main()