1import math, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED=1469
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10    if device.type=='cuda':
 11        torch.zeros(1,device=device)
 12except Exception:
 13    device=torch.device('cpu')
 14
 15# Exact stochastic scalar map: x[t+1]=exp(mu+sigma*eps[t])*x[t].
 16# Its FTLE is exactly Normal(mu, sigma^2/T), making this a direct sanity check.
 17def toy_check():
 18    rng=np.random.default_rng(SEED)
 19    mu=-0.08; z=1.645; n=200000
 20    rows=[]
 21    for T in [4,16,64]:
 22        for sigma in [0.05,0.15,0.30]:
 23            lam=mu + sigma*rng.standard_normal((n,))/math.sqrt(T)
 24            obs_mu=float(lam.mean()); obs_sd=float(lam.std(ddof=1))
 25            p=float((lam>0).mean())
 26            pred_p=0.5*math.erfc((-mu)/(math.sqrt(2)*sigma/math.sqrt(T)))
 27            rows.append({'T':T,'sigma':sigma,'mu_obs':obs_mu,'mu_pred':mu,
 28                         'sd_obs':obs_sd,'sd_pred':sigma/math.sqrt(T),
 29                         'p_obs':p,'p_pred':pred_p})
 30    # Boundary sweep at T=16: UCB changes sign at sigma=-mu*sqrt(T)/z.
 31    T=16; boundary=-mu*math.sqrt(T)/z
 32    boundary_rows=[]
 33    for sigma in [0.12,0.18,0.195,0.23,0.30]:
 34        lam=mu + sigma*rng.standard_normal((n,))/math.sqrt(T)
 35        ucb=float(lam.mean()+z*lam.std(ddof=1))
 36        boundary_rows.append({'sigma':sigma,'ucb_obs':ucb,
 37                              'p_obs':float((lam>0).mean()),
 38                              'ucb_pred':mu+z*sigma/math.sqrt(T)})
 39    return {'predictions':{
 40        'variance_scaling':'sd(lambda_T)=sigma/sqrt(T)',
 41        'tail':'p_+=Phi(mu*sqrt(T)/sigma)',
 42        'boundary':f'ucb=0 at sigma={boundary:.5f} for mu={mu}, T={T}, z={z}'},
 43        'scaling_rows':rows,'boundary_rows':boundary_rows}
 44
 45class NoisyRNN(nn.Module):
 46    def __init__(self,h=12):
 47        super().__init__(); self.h=h
 48        self.W=nn.Parameter(torch.randn(h,h)*0.45)
 49        self.U=nn.Parameter(torch.randn(h,2)*0.35)
 50        self.b=nn.Parameter(torch.zeros(h)); self.out=nn.Linear(h,1)
 51    def run(self,x, noise_sigma=0.0, return_ftle=False, K=4):
 52        B,T,_=x.shape; h=torch.zeros(B,self.h,device=x.device)
 53        if return_ftle:
 54            qs=[torch.randn(B,self.h,device=x.device) for _ in range(K)]
 55            qs=[q/(q.norm(dim=1,keepdim=True)+1e-8) for q in qs]
 56            sums=[torch.zeros(B,device=x.device) for _ in range(K)]
 57        for t in range(T):
 58            eps=torch.randn(K,self.h,self.h,device=x.device) if noise_sigma else None
 59            Wt=self.W + (noise_sigma*eps[0] if eps is not None else 0)
 60            pre=h@Wt.T + x[:,t]@self.U.T + self.b
 61            h=torch.tanh(pre)
 62            if return_ftle:
 63                deriv=1-torch.tanh(pre)**2
 64                new_qs=[]; new_sums=[]
 65                for k in range(K):
 66                    wk=self.W + (noise_sigma*eps[k] if eps is not None else 0)
 67                    q=(qs[k]@wk.T)*deriv
 68                    norm=q.norm(dim=1)+1e-8
 69                    new_qs.append(q/norm[:,None])
 70                    new_sums.append(sums[k]+torch.log(norm))
 71                qs, sums=new_qs, new_sums
 72        y=self.out(h).squeeze(-1)
 73        if return_ftle:
 74            return y, torch.stack(sums,dim=0)/T
 75        return y
 76
 77def make_batch(B=64,T=20):
 78    x=torch.rand(B,T,2,device=device)
 79    # Two marker positions, target is sum of the marked values in channel 0.
 80    p1=torch.randint(0,T//2,(B,),device=device); p2=torch.randint(T//2,T,(B,),device=device)
 81    x[:,:,1]=0
 82    x[torch.arange(B,device=device),p1,1]=1
 83    x[torch.arange(B,device=device),p2,1]=1
 84    y=x[:,:,0].gather(1,p1[:,None]).squeeze(1)+x[:,:,0].gather(1,p2[:,None]).squeeze(1)
 85    return x,y
 86
 87def train_variant(kind, steps=220, T=20, sigma=.08, rho=.35):
 88    torch.manual_seed(SEED+{'base':0,'mean':1,'ucb':2}[kind])
 89    model=NoisyRNN().to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 90    z=1.645; history=[]
 91    for step in range(steps):
 92        x,y=make_batch(64,T)
 93        opt.zero_grad(); pred,lams=model.run(x,sigma,True,K=4)
 94        task=((pred-y)**2).mean(); mu=lams.mean(); sd=lams.std(unbiased=True)
 95        risk=torch.relu(mu + (z*sd if kind=='ucb' else 0.0))**2 if kind!='base' else mu*0
 96        # mean variant penalizes only positive estimated mean FTLE.
 97        loss=task + rho*risk
 98        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
 99        history.append((float(task.detach()),float(mu.detach()),float(sd.detach()),float(loss.detach())))
100    # Clean task and independent noisy FTLE evaluation.
101    with torch.no_grad():
102        vals=[]; ys=[]; ts=[]
103        for _ in range(12):
104            x,y=make_batch(64,T); pred,lams=model.run(x,sigma,True,K=8)
105            vals.append(float(((pred-y)**2).mean())); ts.append(lams.cpu().numpy().ravel())
106        arr=np.concatenate(ts)
107    return {'final_train':history[-1], 'cleanish_mse':float(np.mean(vals)),
108            'ftle_mean':float(arr.mean()),'ftle_sd':float(arr.std(ddof=1)),
109            'positive_fraction':float((arr>0).mean()),
110            'ucb':float(arr.mean()+z*arr.std(ddof=1))}
111
112def main():
113    global device
114    out={'device':str(device),'toy':toy_check(),'rnn':{}}
115    for k in ['base','mean','ucb']:
116        try: out['rnn'][k]=train_variant(k)
117        except Exception as e:
118            if device.type=='cuda':
119                device=torch.device('cpu')
120                out['rnn'][k]={'error':str(e),'fallback':'cpu'}
121            else: raise
122    with open('results.json','w') as f: json.dump(out,f,indent=2)
123    print(json.dumps(out,indent=2))
124if __name__=='__main__': main()