1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import train_model, evaluate, sweep_baseline, make_report
  8from bench.custom_tracks import burgers_periodic_split as track
  9
 10META = {'name':'burgers_front_tokens_bench','domain':'pde','description':'Burgers one-step regression with analytical RH front tokens versus matched CNN.'}
 11
 12# Shared base architecture: both systems have the same hidden convolutional encoder.
 13class Encoder(nn.Module):
 14    def __init__(self, n=32, width=32):
 15        super().__init__(); self.net=nn.Sequential(nn.Conv1d(1,width,5,padding=2),nn.ReLU(),nn.Conv1d(width,width,5,padding=2),nn.ReLU(),nn.Conv1d(width,width,3,padding=1),nn.ReLU())
 16    def forward(self,x): return self.net(x)
 17
 18class Baseline(nn.Module):
 19    def __init__(self,n=32):
 20        super().__init__(); self.enc=Encoder(n); self.head=nn.Conv1d(32,1,1)
 21    def forward(self,x): return self.head(self.enc(x))
 22
 23def front_info(x, threshold=.18):
 24    # Differentiable-free detector is intentionally input preprocessing; traces are noisy local estimates.
 25    u=x[:,0,:]; d=(u[:,:,None] if False else None)
 26    jumps=(u[:,1:]-u[:,:-1]).abs(); B,N=u.shape
 27    mask=jumps > threshold
 28    # one token: strongest detected edge (periodic edge included)
 29    vals, idx=jumps.max(dim=1); idx=idx.clamp(1,N-2)
 30    grid=torch.arange(N,device=x.device).float()[None,:]
 31    pos=(idx.float()+.5)/(N-1)
 32    il=idx-1; ir=idx+1
 33    ul=u.gather(1,il[:,None]).squeeze(1); ur=u.gather(1,ir[:,None]).squeeze(1)
 34    active=(vals>threshold).float()
 35    return pos,ul,ur,active
 36
 37class FrontTokens(nn.Module):
 38    def __init__(self,n=32):
 39        super().__init__(); self.enc=Encoder(n)
 40        # residual head uses same parameter count scale, while traces are estimated from local observed grid values
 41        self.res=nn.Conv1d(32,1,1)
 42        self.trace=nn.Sequential(nn.Linear(32,16),nn.Tanh(),nn.Linear(16,2))
 43        self.n=n
 44    def forward(self,x):
 45        h=self.enc(x); residual=self.res(h)
 46        pos,ul,ur,active=front_info(x)
 47        # RH speed for Burgers: (F(ur)-F(ul))/(ur-ul) = (ul+ur)/2.
 48        speed=.5*(ul+ur)
 49        newpos=(pos + .08*speed).remainder(1.0)
 50        # Render a periodic step from estimated traces; trace MLP receives local encoded context.
 51        b=torch.arange(x.shape[0],device=x.device)
 52        ctx=h[b,:,torch.clamp((pos*(self.n-1)).long(),0,self.n-1)]
 53        delta=self.trace(ctx); aul=ul+0.15*delta[:,0]; aur=ur+0.15*delta[:,1]
 54        xx=torch.arange(self.n,device=x.device).float()[None,:]/(self.n-1)
 55        # smooth sigmoid rendering keeps gradients through positions and traces
 56        dist=((xx-newpos[:,None]+.5)%1.0)-.5
 57        front=aur[:,None]+(aul-aur)[:,None]*torch.sigmoid(-dist/0.035)
 58        return residual + active[:,None,None]*front[:,None,:], (pos,newpos,ul,ur,speed,active)
 59
 60def idea_forward(net,x): return net(x)[0]
 61def ds(seed):
 62    d=track.get_dataset(seed, n_train=400, n_test=160)
 63    for k in ('xtr','ytr','xte','yte'):
 64        d[k]=torch.as_tensor(d[k], dtype=torch.float32)
 65    return d
 66def seed_all(s): random.seed(s); np.random.seed(s); torch.manual_seed(s)
 67def train_one(kind,cfg,seed,return_net=False):
 68    seed_all(seed); d=ds(seed); net=Baseline() if kind=='base' else FrontTokens()
 69    if kind=='base': out=train_model(net,d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *_:None)
 70    else:
 71        # Official trainer expects tensor output; wrap the idea system without changing its learned system.
 72        class Wrap(nn.Module):
 73            def __init__(self,m): super().__init__(); self.m=m
 74            def forward(self,z): return self.m(z)[0]
 75        out=train_model(Wrap(net),d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *_:None)
 76    if return_net: return out,out[0],d
 77    return out[1] if out[1] is not None else float('nan')
 78
 79def signature():
 80    pred=[]; obs=[]
 81    try:
 82        for s in range(8):
 83            out,net,d=train_one('idea',{'lr':.003,'epochs':15},s,True)
 84            if out[0] is None: continue
 85            net.cpu(); net.eval()
 86            x=torch.as_tensor(d['xte'], dtype=torch.float32); y=torch.as_tensor(d['yte'], dtype=torch.float32)
 87            with torch.no_grad(): z,info=net.m(x); pos,np_,ul,ur,sp,act=info
 88            def cross(a):
 89                a=a[:,0,:] if a.ndim==3 else a
 90                q=(a[:,:-1]>=0.5)&(a[:,1:]<0.5); ix=q.float().argmax(1); return (ix.float()+.5)/31
 91            target=cross(y); observed=target-cross(x); m=act>0
 92            if m.any(): pred += ((np_-pos)[m]).tolist(); obs += observed[m].tolist()
 93    except Exception as e:
 94        return {'predicted_mean_speed_displacement':float('nan'),'observed_mean_displacement':float('nan'),'relative_error':float('nan'),'confirmed':False,'error':str(e)[:180]}
 95    if not pred: return {'predicted_mean_speed_displacement':float('nan'),'observed_mean_displacement':float('nan'),'relative_error':float('nan'),'confirmed':False}
 96    a,b=float(np.mean(pred)),float(np.mean(obs)); rel=abs(a-b)/(abs(b)+1e-6)
 97    return {'predicted_mean_speed_displacement':a,'observed_mean_displacement':b,'relative_error':rel,'confirmed':bool(rel<.15)}
 98
 99def main():
100    # Union-parity grid: all idea learning rates are also baseline-evaluated.
101    grid=[{'lr':lr,'epochs':ep} for lr in (.0015,.003,.006) for ep in (15,)]
102    base=sweep_baseline(lambda c: lambda s:train_one('base',c,s),grid)
103    idea_grid=grid
104    idea_tuning=[]
105    for cfg in idea_grid:
106        r=evaluate(lambda s,c=cfg:train_one('idea',c,s), seeds=(0,1,2,3))
107        idea_tuning.append({'cfg':cfg,'mean':r['mean']})
108    idea_cfg=min(idea_tuning,key=lambda z:z['mean'])['cfg']
109    idea=evaluate(lambda s:train_one('idea',idea_cfg,s))
110    idea['selected_cfg']=idea_cfg
111    extra={'idea_sweep':idea_tuning,'custom_track':{'name':'burgers_front_tokens_bench','file':'bench_front_tokens.py','domain':'pde'},'mechanism_signature':signature()}
112    rep=make_report('burgers_periodic_split','cnn_shared',base,idea,extra)
113    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
114if __name__=='__main__': main()