1import json, math, os, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=1729
  7
  8def seed_all(s=SEED):
  9    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 10    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 11
 12class HydroAllocator:
 13    """Conservative path-current allocation of blockwise update activity."""
 14    def __init__(self, m, alpha=0.15, D0=0.7, beta=1.0, noise=0.0, eps=1e-8):
 15        self.m=m; self.alpha=alpha; self.D0=D0; self.beta=beta; self.noise=noise; self.eps=eps
 16        self.n=np.ones(m, dtype=np.float64); self.var=np.zeros(m); self.mean=np.zeros(m)
 17    def step(self, activity, rng):
 18        a=np.asarray(activity, dtype=np.float64)
 19        raw=a + self.eps
 20        target=self.m*raw/raw.sum()
 21        self.mean=0.95*self.mean+0.05*a
 22        self.var=0.95*self.var+0.05*(a-self.mean)**2
 23        # D(n) is positive nonlinear diffusivity; mobility is nonnegative.
 24        nf=.5*(self.n[:-1]+self.n[1:])
 25        D=self.D0*(1.0+self.beta*nf/self.m)
 26        mobility=0.02 + self.var[:-1]+self.var[1:]
 27        z=rng.normal(size=self.m-1)
 28        J=-D*(self.n[1:]-self.n[:-1]) + self.noise*np.sqrt(mobility)*z
 29        div=np.zeros(self.m); div[:-1] += J; div[1:] -= J
 30        self.n=np.maximum(self.n-self.alpha*div, 1e-12)
 31        self.n *= self.m/self.n.sum()
 32        # Activity target is used as the local source only through a gentle blend;
 33        # transport remains the defining conservative mechanism.
 34        self.n=0.85*self.n+0.15*target
 35        self.n *= self.m/self.n.sum()
 36        return self.n.copy()
 37
 38class TinyMLP(nn.Module):
 39    def __init__(self, d=20, h=32, out=3):
 40        super().__init__(); self.layers=nn.ModuleList([nn.Linear(d,h),nn.Linear(h,h),nn.Linear(h,h),nn.Linear(h,out)])
 41    def forward(self,x):
 42        for i,l in enumerate(self.layers): x=l(x); x=torch.tanh(x) if i<3 else x
 43        return x
 44
 45def make_data(n=2400):
 46    g=np.random.default_rng(SEED); x=g.normal(size=(n,20)).astype('float32')
 47    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]])
 48    y=logits.argmax(1).astype('int64'); return torch.tensor(x),torch.tensor(y)
 49
 50def train(mode, device, steps=500):
 51    seed_all(SEED); x,y=make_data(); x=x.to(device); y=y.to(device)
 52    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)]
 53    # Adam moments, with hydrodynamic allocation multiplying the Adam direction.
 54    m1=[torch.zeros_like(p) for p in params]; m2=[torch.zeros_like(p) for p in params]
 55    allocator=HydroAllocator(4, alpha=.18, D0=.8, beta=2., noise=(.015 if mode=='hydro_noise' else 0.0)) if mode!='adam' else None
 56    rng=np.random.default_rng(SEED+11); losses=[]; t0=time.time(); bs=64
 57    for step in range(1,steps+1):
 58        ix=((step-1)*bs)% (len(x)-bs); xb=x[ix:ix+bs]; yb=y[ix:ix+bs]
 59        model.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb),yb); loss.backward()
 60        activities=[]; dirs=[]
 61        with torch.no_grad():
 62            for bi,ps in enumerate(blocks):
 63                sq=0.; norm=0.
 64                for p in ps:
 65                    j=pindex[id(p)]; g=p.grad
 66                    m1[j].mul_(0.9).add_(g,alpha=.1); m2[j].mul_(0.999).addcmul_(g,g,value=.001)
 67                    d=m1[j]/(m2[j].sqrt()+1e-8); dirs.append((j,d))
 68                    sq += float((d*d).sum()); norm += float((p*p).sum())
 69                activities.append(math.sqrt(sq)/(math.sqrt(norm)+1e-8))
 70            if allocator is None: scales=np.ones(4)
 71            else: scales=allocator.step(activities,rng)/(4/4)
 72            # cosine schedule, same base update budget
 73            lr=.012*0.5*(1+math.cos(math.pi*(step-1)/steps))
 74            for bi,ps in enumerate(blocks):
 75                scale=1.0 if allocator is None else float(scales[bi])
 76                for p in ps:
 77                    j=pindex[id(p)]; p.add_(dirs[j][1],alpha=-lr*scale)
 78        losses.append(float(loss.detach().cpu()))
 79    with torch.no_grad(): final=float(nn.functional.cross_entropy(model(x),y).cpu())
 80    return {'final_full_loss':final,'last50_mean':float(np.mean(losses[-50:])),'time_sec':time.time()-t0,
 81            'loss_start':losses[0], 'loss_curve':losses, 'n':None if allocator is None else allocator.n.tolist()}
 82
 83def math_check():
 84    # Deterministic linear path diffusion: n' = n - alpha D L n, zero-flux boundaries.
 85    M=12; D=.73; nbar=1.; rng=np.random.default_rng(SEED)
 86    # edge-current convention gives graph Laplacian with endpoint degree 1.
 87    L=np.diag([1]+[2]*(M-2)+[1])-np.diag(np.ones(M-1),1)-np.diag(np.ones(M-1),-1)
 88    lam=np.linalg.eigvalsh(L); lmax=float(lam[-1]); k=1
 89    v=np.cos(np.pi*k*(np.arange(M)+.5)/M); v-=v.mean(); v/=np.linalg.norm(v)
 90    def run(alpha, nstep=100):
 91        q=v.copy(); amps=[]
 92        for _ in range(nstep): amps.append(abs(q@v)); q=q-alpha*D*(L@q)
 93        return np.array(amps)
 94    a=.5/(D*lam[-1]); amps=run(a); slope=float(np.polyfit(np.arange(1,50),np.log(amps[1:50]),1)[0])
 95    predicted=math.log(1-a*D*lam[k]); stable=run(2.02/(D*lmax),20)
 96    # conservation under nonlinear/noisy allocator
 97    h=HydroAllocator(M,alpha=.15); before=h.n.sum(); sums=[]
 98    for _ in range(100): sums.append(h.step(rng.random(M),rng).sum())
 99    return {'lambda_max':lmax,'alpha_critical_pred':2/(D*lmax),'decay_slope_observed':slope,
100            'decay_slope_predicted':predicted,'slope_ratio':slope/predicted,
101            'mass_error_max':float(np.max(np.abs(np.asarray(sums)-before))),
102            'above_boundary_max_abs':float(np.max(np.abs(stable))),
103            'stable_alpha_used':a,'unstable_alpha_used':2.02/(D*lmax)}
104
105def main():
106    seed_all(); mathres=math_check(); device='cuda' if torch.cuda.is_available() else 'cpu'
107    try:
108        results={m:train(m,device) for m in ['adam','hydro','hydro_noise']}
109    except Exception as e:
110        if device=='cuda':
111            device='cpu'; results={m:train(m,device) for m in ['adam','hydro','hydro_noise']}
112        else: raise
113    out={'device':device,'math':mathres,'training':results}
114    with open('results.json','w') as f: json.dump(out,f,indent=2)
115    print(json.dumps(out,indent=2))
116if __name__=='__main__': main()