1import json, math, time, random
  2import numpy as np
  3import torch
  4
  5SEED = 641
  6random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  7try:
  8    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  9except Exception:
 10    device = torch.device('cpu')
 11
 12# Small normalized 3-D Fourier grid.  P1 is a plane and P2 is a homogeneous cubic
 13# whose restriction to that plane gives a curved/straight algebraic locus.
 14N = 16
 15f = torch.fft.fftfreq(N, d=1.0/N)
 16k1 = f / (N/2)
 17Kz, Ky, Kx = torch.meshgrid(k1, k1, k1, indexing='ij')
 18K = torch.stack([Kx, Ky, Kz], dim=-1).to(device)
 19
 20def poly_residual(K, a, b, cubic):
 21    a = a / (torch.linalg.vector_norm(a) + 1e-8)
 22    cubic = cubic / (torch.linalg.vector_norm(cubic) + 1e-8)
 23    x,y,z = K[...,0], K[...,1], K[...,2]
 24    # ten monomials of a homogeneous cubic
 25    mons = torch.stack([x**3, y**3, z**3, x*x*y, x*x*z, y*y*x, y*y*z,
 26                        z*z*x, z*z*y, x*y*z], -1)
 27    p1 = (K*a).sum(-1)
 28    p2 = (K*b).sum(-1) + (mons*cubic).sum(-1)
 29    return p1*p1 + p2*p2
 30
 31def gate_from_coeff(K, a, b, cubic, tau):
 32    r = poly_residual(K, a, b, cubic)
 33    # Explicit symmetrization is the advertised inversion protection.
 34    rev = poly_residual(-K, a, b, cubic)
 35    r = .5*(r + rev)
 36    return torch.exp(-r / tau), r
 37
 38# Ground-truth curved pinch curve: plane z=0 and cubic x^3-3xy^2=0.
 39a0 = torch.tensor([0.,0.,1.], device=device)
 40b0 = torch.zeros(3, device=device)
 41c0 = torch.zeros(10, device=device); c0[0] = 1.; c0[5] = -3.
 42tau0 = 0.045
 43with torch.no_grad():
 44    target_gate, target_r = gate_from_coeff(K,a0,b0,c0,tau0)
 45
 46# Stage-1 math checks: inversion invariance, and concentration among low residual points.
 47with torch.no_grad():
 48    inv_err = (target_r - poly_residual(-K, a0, b0, c0)).abs().max().item()
 49    flat = target_r.flatten()
 50    q10 = torch.quantile(flat, .10)
 51    q90 = torch.quantile(flat, .90)
 52    gate_flat = target_gate.flatten()
 53    low_gate_mean = gate_flat[flat <= q10].mean().item()
 54    high_gate_mean = gate_flat[flat >= q90].mean().item()
 55    concentration_ratio = low_gate_mean / (high_gate_mean + 1e-12)
 56    active_fraction = (target_gate > .5).float().mean().item()
 57
 58class DenseMultiplier(torch.nn.Module):
 59    def __init__(self, shape):
 60        super().__init__()
 61        self.re = torch.nn.Parameter(torch.zeros(shape))
 62        self.im = torch.nn.Parameter(torch.zeros(shape))
 63        torch.nn.init.normal_(self.re, std=.02)
 64    def forward(self, X):
 65        return X * torch.complex(self.re, self.im)
 66
 67class PinchMultiplier(torch.nn.Module):
 68    def __init__(self):
 69        super().__init__()
 70        self.a = torch.nn.Parameter(torch.randn(3))
 71        self.b = torch.nn.Parameter(torch.randn(3)*.05)
 72        self.cubic = torch.nn.Parameter(torch.randn(10)*.05)
 73        self.log_tau = torch.nn.Parameter(torch.tensor(math.log(.06)))
 74    def forward(self, X):
 75        tau = torch.exp(self.log_tau).clamp(.005, .5)
 76        g, _ = gate_from_coeff(K, self.a, self.b, self.cubic, tau)
 77        return X * g
 78
 79# Random fields make the Fourier regression identify every multiplier entry.
 80Btrain, Bval = 96, 48
 81xtrain = torch.randn(Btrain,N,N,N,device=device)
 82xval = torch.randn(Bval,N,N,N,device=device)
 83with torch.no_grad():
 84    ytrain = torch.fft.ifftn(torch.fft.fftn(xtrain)*target_gate).real
 85    yval = torch.fft.ifftn(torch.fft.fftn(xval)*target_gate).real
 86
 87def fit(model, steps=350):
 88    opt = torch.optim.Adam(model.parameters(), lr=.035)
 89    model.train(); t0=time.perf_counter()
 90    for step in range(steps):
 91        # fixed full batch keeps the comparison deterministic and tiny
 92        pred = torch.fft.ifftn(model(torch.fft.fftn(xtrain))).real
 93        loss = torch.mean((pred-ytrain)**2)
 94        opt.zero_grad(); loss.backward(); opt.step()
 95    elapsed=time.perf_counter()-t0
 96    model.eval()
 97    with torch.no_grad():
 98        pred = torch.fft.ifftn(model(torch.fft.fftn(xval))).real
 99        mse = torch.mean((pred-yval)**2).item()
100        rel = (torch.linalg.vector_norm(pred-yval)/torch.linalg.vector_norm(yval)).item()
101    return mse, rel, elapsed
102
103# CUDA OOM or driver problems are explicitly handled by rerunning on CPU.
104def run():
105    dense = DenseMultiplier((N,N,N)).to(device)
106    pinch = PinchMultiplier().to(device)
107    dmse,drel,dt = fit(dense)
108    pmse, prel, pt = fit(pinch)
109    return dmse,drel,dt,pmse,prel,pt
110try:
111    dmse,drel,dt,pmse,prel,pt = run()
112except Exception as e:
113    device=torch.device('cpu')
114    Kz, Ky, Kx = torch.meshgrid(k1,k1,k1,indexing='ij'); K=torch.stack([Kx,Ky,Kz],-1)
115    target_gate,target_r=gate_from_coeff(K,a0.cpu(),b0.cpu(),c0.cpu(),tau0)
116    xtrain=torch.randn(Btrain,N,N,N); xval=torch.randn(Bval,N,N,N)
117    ytrain=torch.fft.ifftn(torch.fft.fftn(xtrain)*target_gate).real
118    yval=torch.fft.ifftn(torch.fft.fftn(xval)*target_gate).real
119    dmse,drel,dt,pmse,prel,pt=run()
120
121result = {
122 'device': str(device), 'grid': [N,N,N],
123 'math_check': {'max_inversion_residual_error': inv_err,
124                'mean_gate_lowest_residual_decile': low_gate_mean,
125                'mean_gate_highest_residual_decile': high_gate_mean,
126                'low_to_high_gate_ratio': concentration_ratio,
127                'fraction_gate_above_half': active_fraction},
128 'benchmark': {'dense_mse': dmse, 'dense_relative_l2': drel, 'dense_seconds': dt,
129               'pinch_mse': pmse, 'pinch_relative_l2': prel, 'pinch_seconds': pt},
130 'parameters': {'dense_complex_entries': 2*N**3, 'pinch_trainable_scalars': 3+3+10+1},
131 'note': 'Target is exactly an algebraic pinch gate applied to white-noise volumes; this tests representation and optimization on the claimed synthetic phenomenon, not a physical PDE.'
132}
133print(json.dumps(result, indent=2))