1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6from scipy.special import erf
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, evaluate, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = tuple(range(4))
 12EPOCHS = 12
 13BATCH = 128
 14HIDDEN = 64
 15RC = math.sqrt(2.0 / math.pi)
 16# Union of all learning rates used by either side; baseline also sweeps wd.
 17GRID = [
 18    {'lr': 0.001, 'weight_decay': 0.0, 'epochs': EPOCHS},
 19    {'lr': 0.003, 'weight_decay': 0.0, 'epochs': EPOCHS},
 20    {'lr': 0.010, 'weight_decay': 0.0, 'epochs': EPOCHS},
 21    {'lr': 0.003, 'weight_decay': 1e-4, 'epochs': EPOCHS},
 22    {'lr': 0.003, 'weight_decay': 1e-3, 'epochs': EPOCHS},
 23]
 24IDEA_GRID = [
 25    {'lr': 0.001, 'weight_decay': 0.0, 'epochs': EPOCHS, 'R': 0.60},
 26    {'lr': 0.003, 'weight_decay': 0.0, 'epochs': EPOCHS, 'R': 0.90},
 27    {'lr': 0.010, 'weight_decay': 0.0, 'epochs': EPOCHS, 'R': 1.20},
 28]
 29
 30def seed_all(seed):
 31    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 32    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 33
 34def math_check():
 35    def f(m, r): return erf(m / (math.sqrt(2)*r))
 36    rows=[]
 37    for r in [0.6, RC, 1.0]:
 38        eps=1e-6
 39        measured=(f(eps,r)-f(-eps,r))/(2*eps)
 40        rows.append({'R':r, 'observed_slope':measured, 'predicted_slope':RC/r,
 41                     'abs_error':abs(measured-RC/r)})
 42    return {'Rc_predicted':RC, 'rows':rows,
 43            'max_abs_error':max(x['abs_error'] for x in rows)}
 44
 45class MatchedRNN(nn.Module):
 46    """Same trainable input/recurrent/head maps for both systems.
 47    Baseline uses saturated continuous state; idea uses ternary probabilities.
 48    """
 49    def __init__(self, idea=False, R=0.9, beta=2.0, gamma=0.5, J=1.0):
 50        super().__init__()
 51        self.inp=nn.Linear(3,HIDDEN)
 52        self.rec=nn.Linear(HIDDEN,HIDDEN,bias=False)
 53        self.head=nn.Linear(HIDDEN,1)
 54        self.idea=idea; self.R=R; self.beta=beta; self.gamma=gamma; self.J=J
 55        self.last_m=[]
 56    def forward(self,x):
 57        seq=x.view(x.shape[0],-1,3)
 58        h=torch.zeros(x.shape[0],HIDDEN,device=x.device,dtype=x.dtype)
 59        self.last_m=[]
 60        if self.idea:
 61            # p ordering (-1,0,+1), initialized at neutral state.
 62            p=torch.zeros(x.shape[0],HIDDEN,3,device=x.device,dtype=x.dtype)
 63            p[:,:,1]=1.
 64            states=torch.tensor([-1.,0.,1.],device=x.device,dtype=x.dtype)
 65            for t in range(seq.shape[1]):
 66                m=(p*states).sum(-1).mean(-1,keepdim=True)
 67                xi=torch.randn_like(m.expand(-1,HIDDEN))*self.R
 68                u=self.inp(seq[:,t])+self.rec(h)+self.J*m+xi
 69                e=torch.stack([-u, torch.zeros_like(u), u],-1)
 70                logits=-self.beta*e
 71                peq=torch.softmax(logits,dim=-1)
 72                p=(1-self.gamma)*p+self.gamma*peq
 73                h=(p*states).sum(-1)
 74                self.last_m.append(float(h.mean().detach().cpu()))
 75        else:
 76            for t in range(seq.shape[1]):
 77                h=torch.tanh(self.inp(seq[:,t])+self.rec(h))
 78                self.last_m.append(float(h.mean().detach().cpu()))
 79        return self.head(h)
 80
 81def train(cfg, seed, idea=False):
 82    seed_all(seed)
 83    ds=get_dataset('dynamics', seed, n_train=400, n_test=200)
 84    net=MatchedRNN(idea=idea, R=cfg.get('R',0.9)).to('cuda' if torch.cuda.is_available() else 'cpu')
 85    dev=next(net.parameters()).device
 86    xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev)
 87    xte,yte=ds['xte'].to(dev),ds['yte'].to(dev)
 88    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg.get('weight_decay',0.0))
 89    lossf=nn.MSELoss(); hist=[]
 90    try:
 91        for _ in range(cfg['epochs']):
 92            perm=torch.randperm(len(xtr),device=dev); total=0.
 93            for i in range(0,len(xtr),BATCH):
 94                ix=perm[i:i+BATCH]; opt.zero_grad(set_to_none=True)
 95                loss=lossf(net(xtr[ix]),ytr[ix]); loss.backward(); opt.step(); total+=float(loss)*len(ix)
 96            hist.append(total/len(xtr))
 97        net.eval()
 98        with torch.no_grad(): metric=float(lossf(net(xte),yte))
 99        # measured model behavior, not an analytical identity
100        sig={'late_m': list(net.last_m[-4:]), 'mean_abs_late_m': float(np.mean(np.abs(net.last_m[-4:]))) }
101        return {'metric':metric,'history':hist,'seed':seed,'signature':sig}
102    except RuntimeError:
103        # Explicit CPU fallback for shared/fragile CUDA environments.
104        seed_all(seed); net=MatchedRNN(idea=idea,R=cfg.get('R',.9)).cpu()
105        xtr,ytr=ds['xtr'],ds['ytr']; xte,yte=ds['xte'],ds['yte']
106        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg.get('weight_decay',0.0))
107        for _ in range(cfg['epochs']):
108            for i in range(0,len(xtr),BATCH):
109                opt.zero_grad(); loss=lossf(net(xtr[i:i+BATCH]),ytr[i:i+BATCH]); loss.backward(); opt.step()
110        with torch.no_grad(): metric=float(lossf(net(xte),yte))
111        return {'metric':metric,'history':[],'seed':seed,'signature':{'late_m':net.last_m[-4:]}}
112
113def metric_fn(idea, cfg, seed): return train(cfg,seed,idea=idea)['metric']
114
115def main():
116    sanity=math_check(); print('MATH_CHECK',json.dumps(sanity))
117    base=sweep_baseline(lambda cfg: lambda s: metric_fn(False,cfg,s), GRID, seeds=SWEEP_SEEDS)
118    # Idea has three a-priori disorder/lr settings; select on the same 0..3 sweep seeds.
119    tried=[]
120    for cfg in IDEA_GRID:
121        r=evaluate(lambda s,cfg=cfg: metric_fn(True,cfg,s), seeds=SWEEP_SEEDS)
122        tried.append({'cfg':cfg,'mean':r['mean']})
123    best_cfg=min(IDEA_GRID,key=lambda c: next(x['mean'] for x in tried if x['cfg']==c))
124    idea=evaluate(lambda s: metric_fn(True,best_cfg,s), seeds=SEEDS)
125    idea['best_cfg']=best_cfg; idea['sweep']=tried
126    # Re-test trained models for a behavioral signature at NN scale.
127    probe=[]
128    for s in SEEDS:
129        z=train(best_cfg,s,idea=True); probe.append(z['signature']['mean_abs_late_m'])
130    observed=float(np.mean(probe))
131    sig={'prediction':'disorder above Rc suppresses persistent mean-field basin magnetization',
132         'predicted_Rc':RC, 'tested_R':best_cfg['R'],
133         'observed_mean_abs_late_m':observed,
134         'observed_per_seed':probe,
135         'confirmed': bool(best_cfg['R']>RC and observed < 0.5)}
136    report=make_report('dynamics','matched_ternary_rnn',base,idea,extra=sig)
137    report['math_sanity']=sanity
138    report['protocol_notes']='Dynamics selected because the idea targets recurrent stability/basin dependence. Baseline and idea share trainable input, recurrent, and head maps; only state update/disorder differs.'
139    Path('bench_report.json').write_text(json.dumps(report,indent=2))
140    print(json.dumps(report,indent=2))
141if __name__=='__main__': main()