1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 1286
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_default_dtype(torch.float64)
  9device = "cuda" if torch.cuda.is_available() else "cpu"
 10
 11
 12def basis(d, p):
 13    from itertools import combinations
 14    return list(combinations(range(d), p))
 15
 16
 17def wedge_matrix(w, d, p):
 18    B, T = basis(d, p), basis(d, p + 1)
 19    K = np.zeros((len(T), len(B)))
 20    for j, a in enumerate(B):
 21        for r in range(d):
 22            if r in a:
 23                continue
 24            seq = (r,) + a
 25            inversions = sum(seq[u] > seq[v] for u in range(len(seq)) for v in range(u + 1, len(seq)))
 26            K[T.index(tuple(sorted(seq))), j] = ((-1) ** inversions) * w[r]
 27    return K
 28
 29
 30def relu_power(s, m):
 31    return torch.relu(s) ** m / math.factorial(m)
 32
 33
 34class ExactBackbone(nn.Module):
 35    def __init__(self, w, b, k=2):
 36        super().__init__()
 37        self.register_buffer("w", torch.as_tensor(w))
 38        self.register_buffer("b", torch.as_tensor(b))
 39        self.k = k
 40
 41    def features(self, x, m):
 42        return relu_power(x @ self.w.T + self.b, m)
 43
 44    def scalar(self, x, coef):
 45        return self.features(x, self.k) @ coef
 46
 47    def gradient_form(self, x, coef):
 48        # d(sum_i coef_i sigma_k(s_i)) = sum_i coef_i sigma_{k-1}(s_i) ds_i
 49        return self.features(x, self.k - 1) @ (coef[:, None] * self.w)
 50
 51
 52def math_checks():
 53    rng = np.random.default_rng(SEED)
 54    d, n = 3, 7
 55    w = rng.normal(size=(n, d)); w /= np.linalg.norm(w, axis=1, keepdims=True)
 56    max_nil = 0.0
 57    for wi in w:
 58        for p in range(d - 1):
 59            max_nil = max(max_nil, np.max(np.abs(wedge_matrix(wi, d, p + 1) @ wedge_matrix(wi, d, p))))
 60    # Finite-difference verification of d sigma_m(s)=sigma_{m-1}(s) ds away from kink.
 61    x = torch.tensor([[0.37, -0.21, 0.19]], requires_grad=True)
 62    wt = torch.tensor(w[:1]); bt = torch.tensor([0.41]); s = x @ wt.T + bt
 63    m = 3
 64    y = relu_power(s, m).sum(); grad = torch.autograd.grad(y, x)[0].detach().numpy()[0]
 65    rhs = (relu_power(s.detach(), m - 1).numpy()[0, 0] * w[0])
 66    identity_err = float(np.max(np.abs(grad - rhs)))
 67    # Predicted operator scaling: ||K_0|| = ||w||, and composition remains zero.
 68    scales = np.array([0.25, 0.5, 1., 2., 4.])
 69    norms = []
 70    nil_scaled = []
 71    for a in scales:
 72        wi = a * w[0]
 73        k0, k1 = wedge_matrix(wi, d, 0), wedge_matrix(wi, d, 1)
 74        norms.append(np.linalg.norm(k0, 2)); nil_scaled.append(np.max(np.abs(k1 @ k0)))
 75    slope = float(np.polyfit(np.log(scales), np.log(norms), 1)[0])
 76    return {"nilpotence_max_abs": float(max_nil), "derivative_identity_max_abs": identity_err,
 77            "scale_sweep": [{"scale": float(a), "K0_norm": float(v), "K1K0_max": float(z)} for a,v,z in zip(scales,norms,nil_scaled)],
 78            "loglog_scaling_exponent_observed": slope, "predicted_scaling_exponent": 1.0}
 79
 80
 81def make_data(n=512):
 82    x = torch.rand(n, 2, device=device)
 83    # exact potential; target vector field is curl-free
 84    phi = torch.sin(math.pi*x[:,0]) * torch.sin(math.pi*x[:,1])
 85    target = torch.stack((math.pi*torch.cos(math.pi*x[:,0])*torch.sin(math.pi*x[:,1]),
 86                          math.pi*torch.sin(math.pi*x[:,0])*torch.cos(math.pi*x[:,1])), 1)
 87    return x, target
 88
 89
 90def train_compare(steps=700):
 91    rng = np.random.default_rng(SEED)
 92    nfeat = 28
 93    w = rng.normal(size=(nfeat,2)); w /= np.linalg.norm(w,axis=1,keepdims=True)
 94    b = rng.uniform(-0.7,0.7,size=nfeat)
 95    x, target = make_data()
 96    exact = ExactBackbone(w,b,k=2).to(device)
 97    coef = nn.Parameter(torch.zeros(nfeat,device=device)); opt = torch.optim.Adam([coef],lr=0.04)
 98    mlp = nn.Sequential(nn.Linear(2,24),nn.Tanh(),nn.Linear(24,24),nn.Tanh(),nn.Linear(24,2)).to(device)
 99    opt2 = torch.optim.Adam(mlp.parameters(),lr=0.01)
100    for _ in range(steps):
101        opt.zero_grad(); pred=exact.gradient_form(x,coef); loss=((pred-target)**2).mean(); loss.backward(); opt.step()
102        opt2.zero_grad(); pred2=mlp(x); loss2=((pred2-target)**2).mean(); loss2.backward(); opt2.step()
103    with torch.no_grad():
104        p=exact.gradient_form(x,coef); q=mlp(x)
105        exact_mse=float(((p-target)**2).mean()); mlp_mse=float(((q-target)**2).mean())
106    # Autodiff divergence/curl of vector outputs on a smaller probe batch.
107    z=x[:64].detach().requires_grad_(True)
108    def violations(model_kind):
109        if model_kind=="exact": out=exact.gradient_form(z,coef)
110        else: out=mlp(z)
111        grads=[]
112        for j in range(2): grads.append(torch.autograd.grad(out[:,j].sum(),z,retain_graph=True)[0])
113        g0,g1=grads
114        curl=g1[:,0]-g0[:,1]
115        return float(torch.sqrt(torch.mean(curl**2)))
116    return {"exact_params": int(nfeat), "mlp_params": sum(p.numel() for p in mlp.parameters()),
117            "exact_mse": exact_mse, "mlp_mse": mlp_mse,
118            "exact_curl_rms": violations("exact"), "mlp_curl_rms": violations("mlp")}
119
120
121if __name__ == "__main__":
122    try:
123        result={"device":device,"math":math_checks(),"comparison":train_compare()}
124    except Exception as e:
125        if device == "cuda":
126            device="cpu"; torch.cuda.empty_cache(); result={"device":device,"math":math_checks(),"comparison":train_compare()}
127        else:
128            raise
129    with open("results.json","w") as f: json.dump(result,f,indent=2)
130    print(json.dumps(result,indent=2))