1import json
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6SEED=1273
 7N=32
 8DEPTH=16
 9DT=.08
10J=1.0
11
12def device():
13    return 'cuda' if torch.cuda.is_available() else 'cpu'
14
15class PropNet(nn.Module):
16    def __init__(self, mode, gamma=.5, target=.15, alpha=.5):
17        super().__init__(); self.mode=mode; self.gamma0=gamma; self.target=target; self.alpha=alpha
18        # two real fields implement damped coherent propagation; parameter count is identical.
19        self.readout=nn.Sequential(nn.Linear(2*N,16),nn.Tanh(),nn.Linear(16,2))
20    def forward(self,x, return_stats=False):
21        # x: batch,N; q and p are the propagated fields.
22        q=x; p=torch.zeros_like(x); gam=torch.full((x.shape[0],),self.gamma0,device=x.device)
23        gs=[]; rs=[]
24        for _ in range(DEPTH):
25            lapq=torch.roll(q,1,1)+torch.roll(q,-1,1)-2*q
26            lapp=torch.roll(p,1,1)+torch.roll(p,-1,1)-2*p
27            q=q+DT*J*lapp
28            p=p-DT*J*lapq-DT*gam[:,None]*p
29            # fixed controller is detached from the optimization graph.
30            corr=(q[:,:-1]*q[:,1:]).mean(1).abs()
31            var=(q*q).mean(1)+1e-6
32            r=corr/var
33            if self.mode=='adaptive':
34                gam=torch.clamp(gam*torch.exp(self.alpha*(r.detach()-self.target)),.02,16.)
35            gs.append(gam.mean().item()); rs.append(r.mean().item())
36        logits=self.readout(torch.cat([q,p],1))
37        if return_stats: return logits, float(np.mean(gs)), float(np.mean(rs)), float(gam.mean())
38        return logits
39
40def batch(n, dev):
41    # Position of a pulse encodes the class; distractor noise stresses propagation.
42    y=torch.randint(0,2,(n,),device=dev); x=.10*torch.randn(n,N,device=dev)
43    pos=torch.where(y==0, torch.randint(3,N//2,(n,),device=dev), torch.randint(N//2,N-3,(n,),device=dev))
44    x[torch.arange(n,device=dev),pos]=1.
45    return x,y
46
47def run(mode,gamma):
48    torch.manual_seed(SEED); np.random.seed(SEED); dev=device()
49    try:
50        net=PropNet(mode,gamma).to(dev); opt=torch.optim.Adam(net.parameters(),lr=3e-3); lossfn=nn.CrossEntropyLoss()
51        for _ in range(160):
52            x,y=batch(96,dev); loss=lossfn(net(x),y); opt.zero_grad(); loss.backward(); opt.step()
53        with torch.no_grad():
54            x,y=batch(512,dev); logits,g,r,gf=net(x,True); acc=(logits.argmax(1)==y).float().mean().item(); loss=lossfn(logits,y).item()
55        return {'loss':loss,'accuracy':acc,'mean_gamma':g,'mean_ratio':r,'final_gamma':gf,'device':dev}
56    except Exception as e:
57        if dev=='cuda':
58            torch.cuda.empty_cache();
59            # CPU fallback is explicit and reproducible.
60            torch.set_default_device('cpu'); return run(mode,gamma)
61        raise
62
63def main():
64    out={'fixed_ballistic':run('fixed',.5),'fixed_diffusive':run('fixed',4.),'adaptive':run('adaptive',.5)}
65    open('mini_results.json','w').write(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
66if __name__=='__main__': main()