1import json, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=3010
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8rng=np.random.default_rng(SEED)
  9try:
 10    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11    if device.type=='cuda': torch.cuda.set_device(0)
 12except Exception: device=torch.device('cpu')
 13
 14# Bistable scalar dynamics: x=+-1 attractors, x=0 repeller.
 15def dyn(x): return x + 0.18*x*(1-x*x)
 16N=rng.normal(size=(1,16)).astype('float32')
 17def obs(x):
 18    x=np.asarray(x).reshape(-1,1)
 19    q=np.concatenate([x,x*x,np.sin(x),np.cos(x)],1).astype('float32')
 20    y=np.zeros((len(x),20),dtype='float32'); y[:,:4]=q; y[:,4:]=.15*q[:,[0]]@N
 21    return y
 22
 23def samples(n):
 24    x=rng.uniform(-1.45,1.45,n); xn=dyn(x)
 25    return torch.tensor(obs(x)),torch.tensor(obs(xn)),x,xn
 26
 27class WM(nn.Module):
 28    def __init__(self):
 29        super().__init__()
 30        self.e=nn.Sequential(nn.Linear(20,32),nn.Tanh(),nn.Linear(32,1))
 31        self.d=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,20))
 32        self.g=nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,1))
 33    def forward(self,x):
 34        z=self.e(x); return z,self.d(z),self.g(z)
 35
 36def train(reg):
 37    tr=samples(2400); va=samples(1000)
 38    model=WM().to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 39    X,Xn=tr[0].to(device),tr[1].to(device)
 40    for _ in range(700):
 41        z,rec,p=model(X); zn=model.e(Xn)
 42        loss=(rec-X).pow(2).mean()+2*(p-zn).pow(2).mean()
 43        if reg: loss += 2*(p-zn.detach()).pow(2).mean()
 44        opt.zero_grad(); loss.backward(); opt.step()
 45    with torch.no_grad():
 46        V,Vn=va[0].to(device),va[1].to(device)
 47        z,rec,p=model(V); zn=model.e(Vn)
 48        residual=(p-zn).abs().cpu().numpy().ravel()
 49        z0=model.e(V).cpu().numpy().ravel(); zp=model.g(model.e(V)).cpu().numpy().ravel()
 50        recon=(rec-V).pow(2).mean().sqrt().item()
 51    return model, residual,z0,zp,recon
 52
 53def graph_cert(z,zp,eps,nb=28):
 54    lo,hi=float(np.percentile(z,1)),float(np.percentile(z,99))
 55    lo-=.15; hi+=.15; w=(hi-lo)/nb
 56    edges=[set() for _ in range(nb)]
 57    for i in range(nb):
 58        a,b=lo+i*w,lo+(i+1)*w; pts=np.linspace(a,b,9)
 59        # empirical padded image; cells intersecting image interval receive edges
 60        y=np.interp(pts,z,zp) if np.all(np.diff(z)>0) else None
 61        # use model-free nearest validation image samples in each source cell
 62        sel=(z>=a)&(z<=b)
 63        if not np.any(sel): continue
 64        ymin,ymax=zp[sel].min()-eps,zp[sel].max()+eps
 65        js=np.where((lo+(np.arange(nb)+1)*w>=ymin)&(lo+np.arange(nb)*w<=ymax))[0]
 66        edges[i].update(js.tolist())
 67    # Tarjan SCC, then terminal components
 68    idx=0; stack=[]; on=set(); ids=[-1]*nb; low=[0]*nb; comps=[]
 69    def go(v):
 70        nonlocal idx
 71        ids[v]=low[v]=idx; idx+=1; stack.append(v); on.add(v)
 72        for q in edges[v]:
 73            if ids[q]<0: go(q); low[v]=min(low[v],low[q])
 74            elif q in on: low[v]=min(low[v],ids[q])
 75        if low[v]==ids[v]:
 76            c=[]
 77            while True:
 78                q=stack.pop(); on.remove(q); c.append(q)
 79                if q==v: break
 80            comps.append(c)
 81    for v in range(nb):
 82        if ids[v]<0: go(v)
 83    terminal=[]
 84    for c in comps:
 85        if not any(q not in c for v in c for q in edges[v]): terminal.append(c)
 86    # margin to cell boundaries, estimated from padded source images.
 87    margin=float('inf')
 88    for i in range(nb):
 89        if edges[i]:
 90            a,b=lo+i*w,lo+(i+1)*w
 91            vals=zp[(z>=a)&(z<=b)]
 92            if len(vals):
 93                margin=min(margin,float(np.min(np.minimum(np.abs(vals-lo),np.abs(vals-hi)))))
 94    return len(terminal), terminal, w, margin
 95
 96def math_check():
 97    # For scalar perturbations, a padded interval contains every true image iff eps>=residual.
 98    g=np.array([-.8,-.2,.3,.9]); true=g+np.array([.01,-.04,.02,-.03]); r=np.abs(true-g)
 99    return {'max_residual':float(r.max()),'contain_at_eps_max':bool(np.all(np.abs(true-g)<=r.max()+1e-12)), 'contain_below_max':bool(np.all(np.abs(true-g)<=r.max()*.5))}
100
101def main():
102    result={'device':str(device),'math_check':math_check()}
103    for name,reg in [('baseline',False),('semiconjugacy',True)]:
104        model,r,z,zp,recon=train(reg)
105        eps=float(np.max(r)); q=float(np.quantile(r,.95))
106        n,comps,w,margin=graph_cert(z,zp,eps)
107        result[name]={'max_residual':eps,'q95_residual':q,'reconstruction_rmse':recon,'cell_width':w,'estimated_margin':margin,'terminal_sccs':n,'certificate':bool(eps<w)}
108    print(json.dumps(result,indent=2))
109    with open('results.json','w') as f: json.dump(result,f,indent=2)
110if __name__=='__main__': main()