1import json, math, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 1144
  7np.random.seed(SEED); torch.manual_seed(SEED)
  8K = 1.0
  9EPS = 1e-6
 10
 11def reconstruct(x, K=1.0, eps=1e-6):
 12    # x is the proper-velocity spatial coordinate
 13    q = (x*x).sum(dim=-1)
 14    t = torch.sqrt(torch.clamp(q - 1.0/K, min=eps))
 15    return torch.cat([t.unsqueeze(-1), x], dim=-1)
 16
 17def lorentz_residual(z, K=1.0):
 18    # Lorentz convention: -t^2 + ||s||^2 = 1/K
 19    return (z[..., 1:]**2).sum(-1) - z[..., 0]**2 - 1.0/K
 20
 21def math_checks():
 22    out = {}
 23    # Prediction 1: for norms above threshold, reconstruction residual is roundoff.
 24    d = 7
 25    radii = np.array([1.001, 1.01, 1.1, 2., 5.])
 26    x = torch.zeros((len(radii), d), dtype=torch.float64)
 27    x[:, 0] = torch.tensor(radii, dtype=torch.float64)
 28    z = reconstruct(x.double(), K, 1e-30)
 29    residual = lorentz_residual(z.double(), K).abs().numpy()
 30    out['reconstruction_max_abs_residual'] = float(residual.max())
 31    out['reconstruction_residuals'] = residual.tolist()
 32    # Prediction 2: validity transition is r*=1/sqrt(K), estimated by bisection.
 33    threshold = 1/math.sqrt(K)
 34    lo, hi = 0.0, 2.0*threshold
 35    for _ in range(70):
 36        mid=(lo+hi)/2
 37        if mid*mid >= 1/K: hi=mid
 38        else: lo=mid
 39    out['predicted_boundary'] = threshold
 40    out['observed_boundary_bisection'] = hi
 41    # Prediction 3: dt/dr=r/sqrt(r^2-1/K), diverges as delta=r-r* shrinks.
 42    deltas = np.array([1e-1, 1e-2, 1e-3, 1e-4, 1e-5])
 43    r = threshold + deltas
 44    analytic = r/np.sqrt(r*r-1/K)
 45    # central finite differences, with steps small relative to each delta
 46    numeric=[]
 47    for rr, dd in zip(r, deltas):
 48        h=dd*1e-4
 49        f=lambda v: math.sqrt(max(v*v-1/K, 1e-30))
 50        numeric.append((f(rr+h)-f(rr-h))/(2*h))
 51    out['derivative_deltas'] = deltas.tolist()
 52    out['derivative_predicted'] = analytic.tolist()
 53    out['derivative_observed'] = numeric
 54    out['derivative_growth_ratio_predicted'] = float(analytic[-1]/analytic[0])
 55    out['derivative_growth_ratio_observed'] = float(numeric[-1]/numeric[0])
 56    # Sweep violation for the literal epsilon-clamped formula: below boundary is not exact.
 57    rs=np.linspace(.2,1.8,17)
 58    xx=torch.zeros((len(rs),3),dtype=torch.float64); xx[:,0]=torch.tensor(rs)
 59    rr=lorentz_residual(reconstruct(xx.double(),K,EPS),K).numpy()
 60    out['violation_sweep']=[{'r':float(a),'abs_residual':float(abs(b))} for a,b in zip(rs,rr)]
 61    return out
 62
 63class PVNet(nn.Module):
 64    def __init__(self, depth=4, width=32, dim=8):
 65        super().__init__(); self.inp=nn.Linear(2,dim); self.layers=nn.ModuleList([nn.Linear(dim,dim) for _ in range(depth-1)]); self.out=nn.Linear(dim,2)
 66    def forward(self,x):
 67        h=torch.tanh(self.inp(x))
 68        for l in self.layers: h=torch.tanh(l(h))
 69        # PV spatial tensor is h; manifold boundary is only reconstructed at interface.
 70        z=reconstruct(h)
 71        return self.out(z[:,1:]), z
 72
 73class ProjectedLorentzNet(nn.Module):
 74    def __init__(self, depth=4, width=32, dim=8):
 75        super().__init__(); self.inp=nn.Linear(2,dim); self.layers=nn.ModuleList([nn.Linear(dim,dim) for _ in range(depth-1)]); self.out=nn.Linear(dim,2)
 76    def project(self,h):
 77        # Standard radial projection makes every intermediate spatial vector valid.
 78        n=torch.linalg.vector_norm(h,dim=-1,keepdim=True)
 79        return h * torch.clamp(1.0/(n+1e-12), min=1.001)
 80    def forward(self,x):
 81        h=self.project(torch.tanh(self.inp(x)))
 82        for l in self.layers: h=self.project(torch.tanh(l(h)))
 83        z=reconstruct(h)
 84        return self.out(z[:,1:]), z
 85
 86def mini_experiment():
 87    torch.manual_seed(SEED)
 88    n=1024
 89    x=torch.randn(n,2); y=((x[:,0]*x[:,1]>0).long())
 90    tr,va=torch.arange(0,768),torch.arange(768,n)
 91    result={}
 92    for depth in (4,12):
 93        for kind, cls in [('baseline',ProjectedLorentzNet),('pv',PVNet)]:
 94            torch.manual_seed(SEED+depth+(0 if kind=='baseline' else 100))
 95            model=cls(depth=depth); opt=torch.optim.Adam(model.parameters(),lr=2e-3); losses=[]; grad=[]; invalid=0
 96            t0=time.perf_counter()
 97            for step in range(250):
 98                opt.zero_grad(); logits,z=model(x[tr]); loss=nn.functional.cross_entropy(logits,y[tr]); loss.backward()
 99                grad.append(float(torch.nn.utils.clip_grad_norm_(model.parameters(),1e9))); opt.step(); losses.append(float(loss))
100                invalid += int((lorentz_residual(z).abs()>1e-5).sum())
101            elapsed=time.perf_counter()-t0
102            with torch.no_grad():
103                logits,z=model(x[va]); acc=float((logits.argmax(1)==y[va]).float().mean()); v=float(lorentz_residual(z).abs().max())
104            result[f'{kind}_depth{depth}']={'final_train_loss':losses[-1],'val_accuracy':acc,'max_final_residual':v,'invalid_count':invalid,'gradient_std':float(np.std(grad)),'seconds_250_steps':elapsed}
105    return result
106
107if __name__=='__main__':
108    torch.set_num_threads(min(4,torch.get_num_threads()))
109    report={'math_checks':math_checks(),'mini_experiment':mini_experiment()}
110    with open('results.json','w') as f: json.dump(report,f,indent=2)
111    print(json.dumps(report,indent=2))