1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7
  8SEED=2540
  9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11
 12# Toy verification: f_a(x)=a*tanh(x), U=[-alpha,alpha].
 13# Since |tanh(x)| is increasing, exact worst-case output on U is a*tanh(alpha).
 14def toy_sweep():
 15    alpha, eps = 1.0, 0.05
 16    exact_boundary=(alpha-eps)/math.tanh(alpha)
 17    rows=[]
 18    for a in np.linspace(.5,1.35,18):
 19        x=np.linspace(-alpha,alpha,200001)
 20        y=a*np.tanh(x)
 21        violation=float(np.mean(np.abs(y)>alpha-eps))
 22        worst=float(np.max(np.abs(y)))
 23        # reachable interval from a dense initial cloud, with empirical diameter
 24        z=np.linspace(-alpha,alpha,20001)
 25        diam=[]
 26        for _ in range(30):
 27            z=a*np.tanh(z); diam.append(float(z.max()-z.min()))
 28        
 29        # Exact symmetric attractor prediction: zero for a<=1; for a>1 solve x=a*tanh(x).
 30        if a <= 1: predicted_diam=0.0
 31        else:
 32            lo,hi=0.0,5.0
 33            for _ in range(60):
 34                mid=(lo+hi)/2
 35                if a*math.tanh(mid)>mid: lo=mid
 36                else: hi=mid
 37            predicted_diam=2*((lo+hi)/2)
 38        rows.append({'a':float(a),'worst':worst,'violation_rate':violation,'diam0':diam[0],'diam30':diam[-1], 'predicted_attractor_diameter':predicted_diam})
 39    observed=min(r['a'] for r in rows if r['violation_rate']>0)
 40    # Boundary is grid-resolved; predicted threshold is exact.
 41    return {'alpha':alpha,'epsilon':eps,'predicted_boundary':exact_boundary,
 42            'observed_first_violating_grid':observed,'rows':rows,
 43            'boundary_abs_error':abs(observed-exact_boundary)}
 44
 45class RNN(nn.Module):
 46    def __init__(self,h=16):
 47        super().__init__(); self.h=h
 48        self.w_in=nn.Linear(1,h); self.w_h=nn.Linear(h,h); self.out=nn.Linear(h,1)
 49    def transition(self,h,x): return torch.tanh(self.w_in(x)+self.w_h(h))
 50    def forward(self,x,h=None):
 51        if h is None: h=torch.zeros(x.shape[0],self.h,device=x.device)
 52        ys=[]
 53        for t in range(x.shape[1]):
 54            h=self.transition(h,x[:,t,:]); ys.append(self.out(h))
 55        return torch.cat(ys,1).squeeze(-1),h
 56
 57def mackey(n=1800):
 58    x=np.zeros(n+31,dtype=np.float32); x[:31]=1.2
 59    for t in range(30,n+30): x[t+1]=x[t]+.2*x[t-30]/(1+x[t-30]**10)-.1*x[t]
 60    x=x[30:]; x=(x-x.mean())/(x.std()+1e-6)
 61    X=np.stack([x[i:i+40] for i in range(n-40)],0)
 62    Y=np.stack([x[i+1:i+41] for i in range(n-40)],0)
 63    return torch.tensor(X[:,:,None]),torch.tensor(Y)
 64
 65def train(trap_lambda, X, Y, steps=500):
 66    device='cuda' if torch.cuda.is_available() else 'cpu'
 67    try:
 68        model=RNN().to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 69        X,Y=X.to(device),Y.to(device)
 70        alpha,eps=.8,.05
 71        for step in range(steps):
 72            ix=torch.arange((step*32)%len(X),((step*32)%len(X))+32,device=device)%len(X)
 73            xb,yb=X[ix],Y[ix]
 74            pred,_=model(xb); loss=F.mse_loss(pred,yb)
 75            if trap_lambda:
 76                h=torch.empty(256,model.h,device=device).uniform_(-alpha,alpha)
 77                z=torch.zeros(256,1,device=device)
 78                yh=model.transition(h,z)
 79                trap=F.softplus(torch.abs(yh)-alpha+eps).mean()
 80                loss=loss+trap_lambda*trap
 81            opt.zero_grad(); loss.backward(); opt.step()
 82        return model,device
 83    except Exception:
 84        device='cpu'; model=RNN(); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 85        for step in range(steps):
 86            ix=torch.arange((step*32)%len(X),((step*32)%len(X))+32)%len(X); xb,yb=X[ix],Y[ix]
 87            pred,_=model(xb); loss=F.mse_loss(pred,yb)
 88            if trap_lambda:
 89                h=torch.empty(256,model.h).uniform_(-.8,.8); yh=model.transition(h,torch.zeros(256,1))
 90                loss=loss+trap_lambda*F.softplus(torch.abs(yh)-.75).mean()
 91            opt.zero_grad(); loss.backward(); opt.step()
 92        return model,device
 93
 94def eval_model(model,device,X,Y,particles=10000,horizon=1000):
 95    alpha,eps=.8,.05
 96    model.eval()
 97    with torch.no_grad():
 98        pred,_=model(X[:128].to(device)); mse=float(F.mse_loss(pred,Y[:128].to(device)).cpu())
 99        h=torch.empty(particles,model.h,device=device).uniform_(-alpha,alpha); zero=torch.zeros(particles,1,device=device)
100        violations=[]; maxnorm=[]; clouds=[]
101        for t in range(horizon):
102            h=model.transition(h,zero); violations.append(float((h.abs()>alpha-eps).any(1).float().mean().cpu()))
103            if t in tuple(sorted(set([0, min(9,horizon-1), min(99,horizon-1), horizon-1]))): clouds.append(float((h.max(0).values-h.min(0).values).norm().cpu()))
104            maxnorm.append(float(h.norm(dim=1).max().cpu()))
105        # one-step empirical certificate on fresh U samples
106        cert=float(np.mean(violations[:1]))
107        return {'mse':mse,'one_step_violation':cert,'max_norm':max(maxnorm),
108                'mean_violation_1000':float(np.mean(violations)),'cloud_diameters':clouds}
109
110def main():
111    toy=toy_sweep(); X,Y=mackey()
112    
113    sweep={}
114    for lam in [0., .1, .3, 1., 3., 10.]:
115        torch.manual_seed(SEED + int(lam*10))
116        m,d=train(lam,X,Y,steps=350)
117        sweep[str(lam)]=eval_model(m,d,X,Y,particles=3000,horizon=300)
118    results={'toy':toy,'rnn_lambda_sweep':sweep}
119    Path('results.json').write_text(json.dumps(results,indent=2))
120    print(json.dumps(results,indent=2))
121if __name__=='__main__': main()