1import json, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5from sklearn.datasets import load_digits
  6from sklearn.model_selection import train_test_split
  7
  8SEED = 530
  9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 10if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)
 11
 12
 13def grid_graph(h, w):
 14    n = h*w
 15    edges = []
 16    for r in range(h):
 17        for c in range(w):
 18            i = r*w+c
 19            if r+1 < h: edges.append((i, (r+1)*w+c))
 20            if c+1 < w: edges.append((i, r*w+c+1))
 21    B = np.zeros((n, len(edges)), dtype=np.float64)
 22    for e, (i, j) in enumerate(edges):
 23        B[i, e] = 1.0; B[j, e] = -1.0
 24    return B, B @ B.T, edges
 25
 26
 27def positive_mass_projection(x, mass):
 28    x = np.maximum(x, 0.0)
 29    s = x.sum(axis=1, keepdims=True)
 30    out = x / np.maximum(s, 1e-12) * mass
 31    dead = (s[:, 0] <= 1e-12)
 32    out[dead] = mass / x.shape[1]
 33    return out
 34
 35
 36def density_step(rho, B, L, edges, dt, D, rng):
 37    # Finite-volume edge density and divergence-form stochastic flux.
 38    edge_rho = np.stack([(rho[:, i] + rho[:, j]) / 2.0 for i, j in edges], axis=1)
 39    xi = rng.standard_normal((rho.shape[0], len(edges)))
 40    stochastic_flux = np.sqrt(np.maximum(edge_rho, 0.0)) * xi
 41    nxt = rho - dt * D * (rho @ L.T) + np.sqrt(2.0 * D * dt) * (stochastic_flux @ B.T)
 42    return positive_mass_projection(nxt, rho.shape[1])
 43
 44
 45def math_check():
 46    B, L, edges = grid_graph(8, 8)
 47    eig = np.linalg.eigvalsh(L)
 48    rng = np.random.default_rng(SEED)
 49    rho = np.ones((4096, 64))
 50    dt = 0.20 / eig.max()  # D dt lambda_max = 0.2, safely below 1
 51    sums=[]; mins=[]; snapshots=[]
 52    for t in range(80):
 53        rho = density_step(rho, B, L, edges, dt, 1.0, rng)
 54        sums.append(np.max(np.abs(rho.sum(1)-64.0)))
 55        mins.append(rho.min())
 56        if t >= 40: snapshots.append(rho.copy())
 57    x = np.concatenate(snapshots, axis=0) - 1.0
 58    # Spatial covariance by Manhattan distance, compared to iid Gaussian with same variance.
 59    coords=[(i//8,i%8) for i in range(64)]
 60    bins={d:[] for d in range(1,9)}
 61    for i,(ri,ci) in enumerate(coords):
 62        for j,(rj,cj) in enumerate(coords):
 63            d=abs(ri-rj)+abs(ci-cj)
 64            if 1 <= d <= 8: bins[d].append((i,j))
 65    cov={d: float(np.mean([np.mean(x[:,i]*x[:,j]) for i,j in pairs])) for d,pairs in bins.items()}
 66    var=float(np.mean(x*x))
 67    iid={d: (0.0 if d>0 else var) for d in bins}
 68    return {'lambda_max':float(eig.max()), 'dt':float(dt),
 69            'max_mass_error':float(max(sums)), 'minimum_density':float(min(mins)),
 70            'variance':var, 'covariance_by_manhattan_distance':cov,
 71            'iid_control_covariance_by_distance':iid,
 72            'mass_and_positivity_pass': bool(max(sums)<1e-10 and min(mins)>=-1e-12)}
 73
 74
 75class TokenNet(nn.Module):
 76    def __init__(self, mode, n=64, dim=16, alpha=0.30, seed=SEED):
 77        super().__init__(); self.mode=mode; self.n=n; self.dim=dim; self.alpha=alpha
 78        self.embed=nn.Linear(1, dim); self.fc=nn.Sequential(nn.LayerNorm(n*dim), nn.Linear(n*dim,64), nn.GELU(), nn.Linear(64,10))
 79        B,L,edges=grid_graph(8,8); self.B=torch.tensor(B,dtype=torch.float32); self.L=torch.tensor(L,dtype=torch.float32); self.edges=edges
 80        self.rng=np.random.default_rng(seed+17)
 81    def forward(self, x):
 82        z=self.embed(x.reshape(x.shape[0], self.n, 1))
 83        if self.training and self.mode != 'none':
 84            if self.mode == 'iid':
 85                noise=torch.randn_like(z)
 86                z=z + self.alpha*noise
 87            else:
 88                b=x.shape[0]; rho=np.ones((b,self.n), dtype=np.float64)
 89                # Autonomous density is deliberately detached from model autograd.
 90                rho=density_step(rho, self.B.numpy(), self.L.numpy(), self.edges, 0.20/float(torch.linalg.eigvalsh(self.L).max()), 1.0, self.rng)
 91                amplitude=torch.tensor((rho-1.0)/np.sqrt(0.65),dtype=z.dtype,device=z.device).unsqueeze(-1)
 92                z=z + self.alpha*amplitude*torch.randn_like(z)
 93        return self.fc(z.flatten(1))
 94
 95
 96def train_eval(mode, Xtr, ytr, Xte, yte, device):
 97    torch.manual_seed(SEED); model=TokenNet(mode).to(device); opt=torch.optim.AdamW(model.parameters(),lr=2e-3,weight_decay=1e-4)
 98    bs=128; losses=[]; t0=time.time()
 99    for epoch in range(15):
100        model.train(); perm=torch.randperm(len(Xtr),device=device); total=0.
101        for st in range(0,len(Xtr),bs):
102            ix=perm[st:st+bs]; out=model(Xtr[ix]); loss=nn.functional.cross_entropy(out,ytr[ix]); opt.zero_grad(); loss.backward(); opt.step(); total += loss.item()*len(ix)
103        losses.append(total/len(Xtr))
104    model.eval()
105    with torch.no_grad(): pred=model(Xte).argmax(1); acc=float((pred==yte).float().mean())
106    return {'test_accuracy':acc,'final_train_loss':losses[-1],'seconds':time.time()-t0,'loss_curve':losses}
107
108
109def main():
110    check=math_check()
111    d=load_digits(); X=d.images.astype('float32')/16.0; y=d.target.astype('int64')
112    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
113    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
114    try:
115        Xtr=torch.tensor(Xtr).to(device); Xte=torch.tensor(Xte).to(device); ytr=torch.tensor(ytr).to(device); yte=torch.tensor(yte).to(device)
116        results={m:train_eval(m,Xtr,ytr,Xte,yte,device) for m in ('none','iid','conserved')}
117    except Exception as e:
118        device=torch.device('cpu'); Xtr=torch.tensor(Xtr).cpu(); Xte=torch.tensor(Xte).cpu(); ytr=torch.tensor(ytr).cpu(); yte=torch.tensor(yte).cpu()
119        results={m:train_eval(m,Xtr,ytr,Xte,yte,device) for m in ('none','iid','conserved')}; results['device_fallback']=str(e)
120    out={'seed':SEED,'device':str(device),'math_check':check,'training':results}
121    with open('results.json','w') as f: json.dump(out,f,indent=2)
122    print(json.dumps(out,indent=2))
123
124if __name__=='__main__': main()