1import json, math, os
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=1234
  7rng=np.random.default_rng(SEED)
  8torch.manual_seed(SEED)
  9try:
 10    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11    if device.type=='cuda':
 12        torch.empty(1, device=device)
 13except Exception:
 14    device=torch.device('cpu')
 15
 16
 17def random_rotations(n, rng):
 18    # Haar rotations from QR, with proper determinant
 19    a=rng.normal(size=(n,3,3)); out=[]
 20    for x in a:
 21        q,r=np.linalg.qr(x); q=q@np.diag(np.sign(np.diag(r)))
 22        if np.linalg.det(q)<0: q[:,0]*=-1
 23        out.append(q)
 24    return np.asarray(out)
 25
 26def lattice_offsets(shells):
 27    # shells are squared Euclidean distances on a cubic periodic lattice
 28    alloff=[]
 29    for d in range(1,shells+1):
 30        off=[]
 31        for x in range(-d,d+1):
 32            for y in range(-d,d+1):
 33                for z in range(-d,d+1):
 34                    if (x,y,z)!=(0,0,0) and x*x+y*y+z*z==d*d:
 35                        off.append((x,y,z))
 36        alloff.append(off)
 37    return alloff
 38
 39def p2(x): return 0.5*(3*x*x-1)
 40
 41def corr_features(u,L,shells=3):
 42    # u: [N,3], row-major cubic lattice
 43    a=u.reshape(L,L,L,3); chans=[]
 44    for offsets in lattice_offsets(shells):
 45        c=np.zeros((L,L,L))
 46        for dx,dy,dz in offsets:
 47            shifted=np.roll(a,(dx,dy,dz),(0,1,2))
 48            c += p2(np.sum(a*shifted,axis=-1))
 49        c /= len(offsets)
 50        chans.append(c)
 51    return np.stack(chans,axis=0).astype(np.float32).reshape(-1)
 52
 53def S_order(u):
 54    Q=np.einsum('ni,nj->ij',u,u)/len(u)
 55    Q=0.5*(3*Q-np.eye(3))
 56    return float(np.linalg.eigvalsh(Q)[-1])
 57
 58def make_config(L, phase, rng):
 59    n=L**3
 60    if phase=='iso':
 61        u=rng.normal(size=(n,3)); u/=np.linalg.norm(u,axis=1,keepdims=True)
 62    else:
 63        d=rng.normal(size=3); d/=np.linalg.norm(d)
 64        # Gaussian tangent noise gives a controllable nematic cloud
 65        u=d+rng.normal(scale=.23,size=(n,3)); u/=np.linalg.norm(u,axis=1,keepdims=True)
 66        # random apolar signs make the physical representation genuinely apolar
 67        u*=rng.choice([-1.,1.],size=(n,1))
 68    R=random_rotations(1,rng)[0]
 69    return u@R.T
 70
 71def transform(u,R,flip_prob,rng):
 72    v=u@R.T
 73    signs=np.where(rng.random(len(u))<flip_prob,-1.,1.)[:,None]
 74    return v*signs
 75
 76class AE(nn.Module):
 77    def __init__(self, dim, zdim=2):
 78        super().__init__(); self.enc=nn.Sequential(nn.Linear(dim,64),nn.Tanh(),nn.Linear(64,zdim)); self.dec=nn.Sequential(nn.Linear(zdim,64),nn.Tanh(),nn.Linear(64,dim))
 79    def forward(self,x):
 80        z=self.enc(x); return self.dec(z),z
 81
 82def train_ae(X, epochs=100):
 83    model=AE(X.shape[1]).to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 84    xt=torch.tensor(X,device=device)
 85    for _ in range(epochs):
 86        opt.zero_grad(); rec,_=model(xt); loss=((rec-xt)**2).mean(); loss.backward(); opt.step()
 87    with torch.no_grad():
 88        rec_t,z_t=model(xt); z=z_t.cpu().numpy(); rec=rec_t.cpu().numpy()
 89    return float(np.mean((rec-X)**2)), z
 90
 91def abs_corr(z,s):
 92    return [float(abs(np.corrcoef(z[:,k],s)[0,1])) for k in range(z.shape[1])]
 93
 94def main():
 95    L=5; n=2*120; us=[]; labels=[]; ss=[]
 96    for phase in ['iso','nematic']:
 97        for _ in range(120):
 98            u=make_config(L,phase,rng); us.append(u); labels.append(phase=='nematic'); ss.append(S_order(u))
 99    us=np.asarray(us); labels=np.asarray(labels); ss=np.asarray(ss)
100    # Core prediction 1: common rotations and independent apolar flips leave c exactly unchanged.
101    base=us[0]; inv=[]; raw=[]
102    for p in [0,.25,.5,.75,1.0]:
103        errs=[]; rawerrs=[]
104        for _ in range(20):
105            v=transform(base,random_rotations(1,rng)[0],p,rng)
106            errs.append(np.linalg.norm(corr_features(base,L,3)-corr_features(v,L,3))/math.sqrt(len(corr_features(base,L,3))))
107            rawerrs.append(np.linalg.norm(base-v)/math.sqrt(base.size))
108        inv.append({'flip_probability':p,'observed_rms_feature_error':float(np.mean(errs)),'raw_rms_vector_error':float(np.mean(rawerrs))})
109    # Core prediction 2: invariant error is independent of shell count (all exact pairwise dot products).
110    shell=[]
111    for R in [1,2,3,4]:
112        vals=[]
113        for _ in range(20):
114            v=transform(base,random_rotations(1,rng)[0],rng.random(),rng)
115            a=corr_features(base,L,R); b=corr_features(v,L,R)
116            vals.append(np.max(np.abs(a-b)))
117        shell.append({'shells':R,'observed_max_abs_error':float(max(vals)),'predicted':0.0})
118    # A direct P2 sweep quantifies evenness: sign of dot product never matters.
119    x=np.linspace(-1,1,10001); even_err=float(np.max(np.abs(p2(x)-p2(-x))))
120    # Mini phase experiment: raw vectors versus invariant local channels.
121    rawX=us.reshape(len(us),-1).astype(np.float32)
122    invX=np.asarray([corr_features(u,L,3) for u in us])
123    raw_mse,zraw=train_ae(rawX); inv_mse,zinv=train_ae(invX)
124    result={'seed':SEED,'device':str(device),'L':L,'n_samples':len(us),
125      'math_predictions':{
126       'rotation_and_sign_invariance':'feature RMS error should be numerical zero for every flip probability',
127       'shell_invariance':'max absolute error should remain numerical zero as shells increase',
128       'P2_evenness':'P2(x)-P2(-x) should be exactly zero'},
129      'invariance_sweep':inv,'shell_sweep':shell,'p2_even_max_error':even_err,
130      'phase_experiment':{'raw_reconstruction_mse':raw_mse,'invariant_reconstruction_mse':inv_mse,
131        'raw_abs_latent_S_correlations':abs_corr(zraw,ss),'invariant_abs_latent_S_correlations':abs_corr(zinv,ss),
132        'raw_best_abs_corr':max(abs_corr(zraw,ss)),'invariant_best_abs_corr':max(abs_corr(zinv,ss)),
133        'mean_S_iso':float(ss[~labels].mean()),'mean_S_nematic':float(ss[labels].mean())}}
134    with open('results.json','w') as f: json.dump(result,f,indent=2)
135    print(json.dumps(result,indent=2))
136
137if __name__=='__main__': main()