import json, math, os, random, time import numpy as np import torch from torch import nn SEED=1729 def seed_all(s=SEED): random.seed(s); np.random.seed(s); torch.manual_seed(s) if torch.cuda.is_available(): torch.cuda.manual_seed_all(s) class HydroAllocator: """Conservative path-current allocation of blockwise update activity.""" def __init__(self, m, alpha=0.15, D0=0.7, beta=1.0, noise=0.0, eps=1e-8): self.m=m; self.alpha=alpha; self.D0=D0; self.beta=beta; self.noise=noise; self.eps=eps self.n=np.ones(m, dtype=np.float64); self.var=np.zeros(m); self.mean=np.zeros(m) def step(self, activity, rng): a=np.asarray(activity, dtype=np.float64) raw=a + self.eps target=self.m*raw/raw.sum() self.mean=0.95*self.mean+0.05*a self.var=0.95*self.var+0.05*(a-self.mean)**2 # D(n) is positive nonlinear diffusivity; mobility is nonnegative. nf=.5*(self.n[:-1]+self.n[1:]) D=self.D0*(1.0+self.beta*nf/self.m) mobility=0.02 + self.var[:-1]+self.var[1:] z=rng.normal(size=self.m-1) J=-D*(self.n[1:]-self.n[:-1]) + self.noise*np.sqrt(mobility)*z div=np.zeros(self.m); div[:-1] += J; div[1:] -= J self.n=np.maximum(self.n-self.alpha*div, 1e-12) self.n *= self.m/self.n.sum() # Activity target is used as the local source only through a gentle blend; # transport remains the defining conservative mechanism. self.n=0.85*self.n+0.15*target self.n *= self.m/self.n.sum() return self.n.copy() class TinyMLP(nn.Module): def __init__(self, d=20, h=32, out=3): super().__init__(); self.layers=nn.ModuleList([nn.Linear(d,h),nn.Linear(h,h),nn.Linear(h,h),nn.Linear(h,out)]) def forward(self,x): for i,l in enumerate(self.layers): x=l(x); x=torch.tanh(x) if i<3 else x return x def make_data(n=2400): g=np.random.default_rng(SEED); x=g.normal(size=(n,20)).astype('float32') w=g.normal(size=(20,3)); logits=x@w + .5*np.sin(x[:,:3])@np.array([[1.,-1.,.5],[.5,.5,-1.],[1.,.2,.3]]) y=logits.argmax(1).astype('int64'); return torch.tensor(x),torch.tensor(y) def train(mode, device, steps=500): seed_all(SEED); x,y=make_data(); x=x.to(device); y=y.to(device) model=TinyMLP().to(device); params=list(model.parameters()); pindex={id(p): j for j,p in enumerate(params)}; blocks=[list(model.layers[i].parameters()) for i in range(4)] # Adam moments, with hydrodynamic allocation multiplying the Adam direction. m1=[torch.zeros_like(p) for p in params]; m2=[torch.zeros_like(p) for p in params] allocator=HydroAllocator(4, alpha=.18, D0=.8, beta=2., noise=(.015 if mode=='hydro_noise' else 0.0)) if mode!='adam' else None rng=np.random.default_rng(SEED+11); losses=[]; t0=time.time(); bs=64 for step in range(1,steps+1): ix=((step-1)*bs)% (len(x)-bs); xb=x[ix:ix+bs]; yb=y[ix:ix+bs] model.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb),yb); loss.backward() activities=[]; dirs=[] with torch.no_grad(): for bi,ps in enumerate(blocks): sq=0.; norm=0. for p in ps: j=pindex[id(p)]; g=p.grad m1[j].mul_(0.9).add_(g,alpha=.1); m2[j].mul_(0.999).addcmul_(g,g,value=.001) d=m1[j]/(m2[j].sqrt()+1e-8); dirs.append((j,d)) sq += float((d*d).sum()); norm += float((p*p).sum()) activities.append(math.sqrt(sq)/(math.sqrt(norm)+1e-8)) if allocator is None: scales=np.ones(4) else: scales=allocator.step(activities,rng)/(4/4) # cosine schedule, same base update budget lr=.012*0.5*(1+math.cos(math.pi*(step-1)/steps)) for bi,ps in enumerate(blocks): scale=1.0 if allocator is None else float(scales[bi]) for p in ps: j=pindex[id(p)]; p.add_(dirs[j][1],alpha=-lr*scale) losses.append(float(loss.detach().cpu())) with torch.no_grad(): final=float(nn.functional.cross_entropy(model(x),y).cpu()) return {'final_full_loss':final,'last50_mean':float(np.mean(losses[-50:])),'time_sec':time.time()-t0, 'loss_start':losses[0], 'loss_curve':losses, 'n':None if allocator is None else allocator.n.tolist()} def math_check(): # Deterministic linear path diffusion: n' = n - alpha D L n, zero-flux boundaries. M=12; D=.73; nbar=1.; rng=np.random.default_rng(SEED) # edge-current convention gives graph Laplacian with endpoint degree 1. L=np.diag([1]+[2]*(M-2)+[1])-np.diag(np.ones(M-1),1)-np.diag(np.ones(M-1),-1) lam=np.linalg.eigvalsh(L); lmax=float(lam[-1]); k=1 v=np.cos(np.pi*k*(np.arange(M)+.5)/M); v-=v.mean(); v/=np.linalg.norm(v) def run(alpha, nstep=100): q=v.copy(); amps=[] for _ in range(nstep): amps.append(abs(q@v)); q=q-alpha*D*(L@q) return np.array(amps) a=.5/(D*lam[-1]); amps=run(a); slope=float(np.polyfit(np.arange(1,50),np.log(amps[1:50]),1)[0]) predicted=math.log(1-a*D*lam[k]); stable=run(2.02/(D*lmax),20) # conservation under nonlinear/noisy allocator h=HydroAllocator(M,alpha=.15); before=h.n.sum(); sums=[] for _ in range(100): sums.append(h.step(rng.random(M),rng).sum()) return {'lambda_max':lmax,'alpha_critical_pred':2/(D*lmax),'decay_slope_observed':slope, 'decay_slope_predicted':predicted,'slope_ratio':slope/predicted, 'mass_error_max':float(np.max(np.abs(np.asarray(sums)-before))), 'above_boundary_max_abs':float(np.max(np.abs(stable))), 'stable_alpha_used':a,'unstable_alpha_used':2.02/(D*lmax)} def main(): seed_all(); mathres=math_check(); device='cuda' if torch.cuda.is_available() else 'cpu' try: results={m:train(m,device) for m in ['adam','hydro','hydro_noise']} except Exception as e: if device=='cuda': device='cpu'; results={m:train(m,device) for m in ['adam','hydro','hydro_noise']} else: raise out={'device':device,'math':mathres,'training':results} with open('results.json','w') as f: json.dump(out,f,indent=2) print(json.dumps(out,indent=2)) if __name__=='__main__': main()