1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED=2380
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device='cuda' if torch.cuda.is_available() else 'cpu'
 11try:
 12    if device=='cuda': torch.cuda.get_device_name(0)
 13except Exception:
 14    device='cpu'
 15
 16# Smooth operators from the proposal.
 17def softmin(a,tau,dim=-1):
 18    return -tau*torch.logsumexp(-a/tau,dim=dim)
 19def softmax_reduce(a,tau,dim=-1):
 20    return tau*torch.logsumexp(a/tau,dim=dim)
 21
 22# First, verify the finite-set approximation claim numerically.
 23def math_check():
 24    rng=np.random.default_rng(SEED)
 25    rows=[]
 26    max_softmin_bound=0.; max_softmax_bound=0.
 27    for n in [2,4,8,16,32]:
 28        a=torch.tensor(rng.normal(size=(500,n)),dtype=torch.float32)
 29        amin=a.min(1).values; amax=a.max(1).values
 30        for tau in [0.2,0.1,0.05,0.02]:
 31            sm=softmin(a,tau); sx=softmax_reduce(a,tau)
 32            # softmin is below min by at most tau log n; softmax above max likewise.
 33            e1=(amin-sm).max().item(); e2=(sx-amax).max().item()
 34            b=tau*math.log(n)
 35            max_softmin_bound=max(max_softmin_bound,e1/b if b else 0)
 36            max_softmax_bound=max(max_softmax_bound,e2/b if b else 0)
 37            rows.append({'n':n,'tau':tau,'softmin_error':float((amin-sm).mean()),
 38                         'softmax_error':float((sx-amax).mean()),'bound':b,
 39                         'min_bound_ratio':e1/b,'max_bound_ratio':e2/b})
 40    # Direct monotone scaling check for one fixed candidate set.
 41    fixed=torch.tensor([[0.1,0.7,1.4,-0.2]],dtype=torch.float32)
 42    taus=torch.tensor([.4,.2,.1,.05,.025])
 43    errs=[float((fixed.min()-softmin(fixed,float(t))).abs()) for t in taus]
 44    return {'rows':rows,'max_bound_ratio_min':max_softmin_bound,
 45            'max_bound_ratio_max':max_softmax_bound,'fixed_tau':taus.tolist(),
 46            'fixed_softmin_abs_error':errs}
 47
 48class Dynamics(nn.Module):
 49    def __init__(self):
 50        super().__init__()
 51        self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
 52    def forward(self,x): return self.net(x)
 53
 54def make_data(seed):
 55    g=torch.Generator().manual_seed(seed)
 56    x=-1+2*torch.rand(256,1,generator=g)
 57    # Smooth, mildly nonlinear concrete map; noisy observations make abstraction useful.
 58    y=.72*x+.18*torch.sin(3*x)+.025*torch.randn(x.shape,generator=g)
 59    return x,y
 60
 61def train(seed,beta,steps=500):
 62    torch.manual_seed(seed)
 63    x,y=make_data(seed+100)
 64    # cells are learnable, but assignments and candidate successor lists are fixed for a fair comparison.
 65    K=16
 66    centers0=torch.linspace(-1.15,1.15,K).view(-1,1)
 67    centers=nn.Parameter(centers0+0.015*torch.randn_like(centers0))
 68    model=Dynamics()
 69    opt=torch.optim.Adam(list(model.parameters())+[centers],lr=3e-3)
 70    # Fixed local successor topology: three cells around the nominal target map.
 71    with torch.no_grad():
 72        nominal=.72*centers0+.18*torch.sin(3*centers0)
 73        nearest=torch.cdist(nominal,centers0).topk(3,largest=False).indices
 74        assign=torch.cdist(x,centers0).argmin(1)
 75    tau=.06
 76    for step in range(steps):
 77        pred=model(x)
 78        mse=((pred-y)**2).mean()
 79        ci=centers[assign]
 80        dcur=(x-ci).abs().squeeze(1)
 81        succ=centers[nearest[assign]].reshape(x.shape[0],-1) # B,3
 82        dnext=(pred-succ).abs()
 83        # Smooth over successors then samples, as in the proposal.
 84        sim=softmax_reduce(softmin(dnext,tau,1)-dcur,tau,0)
 85        violation=torch.relu(dcur-.10)
 86        eps_loss=softmax_reduce(violation,tau,0)
 87        loss=mse+beta*(sim+0.15*eps_loss)
 88        opt.zero_grad(); loss.backward(); opt.step()
 89        with torch.no_grad(): centers.clamp_(-1.3,1.3)
 90    with torch.no_grad():
 91        pred=model(x); ci=centers[assign]; cur=(x-ci).abs().squeeze(1)
 92        succ=centers[nearest[assign]].reshape(x.shape[0],-1); dn=(pred-succ).abs()
 93        # Independent hard reverse-simulation residual and coverage at epsilon=.10.
 94        hard=(dn.min(1).values-cur).max().item()
 95        hard_mean=(dn.min(1).values-cur).mean().item()
 96        covered=((dn.min(1).values<=.10)&(cur<=.10)).float().mean().item()
 97        pred_mse=((pred-y)**2).mean().item()
 98        # Number of candidate successors needed for each sample at radius epsilon.
 99        counts=(dn<=.10).sum(1).float()
100        avg_count=counts.mean().item(); zero=(counts==0).float().mean().item()
101        sim_value=float(softmax_reduce(softmin(dn,tau,1)-cur,tau,0))
102    return {'seed':seed,'beta':beta,'prediction_mse':pred_mse,'hard_max_residual':hard,
103            'hard_mean_residual':hard_mean,'coverage_at_eps':covered,'avg_successors_at_eps':avg_count,
104            'zero_coverage_fraction':zero,'smooth_sim':sim_value}
105
106def main():
107    check=math_check()
108    results=[]
109    for beta in [0.,0.15,0.5]:
110        for seed in [11,22,33]:
111            try: results.append(train(seed,beta))
112            except RuntimeError as e:
113                if 'CUDA' in str(e) or 'out of memory' in str(e).lower():
114                    global device; device='cpu'; results.append(train(seed,beta))
115                else: raise
116    out={'device':device,'math_check':check,'results':results}
117    Path('results.json').write_text(json.dumps(out,indent=2))
118    for b in [0.,.15,.5]:
119        r=[z for z in results if z['beta']==b]
120        vals=np.mean([[q['prediction_mse'],q['hard_max_residual'],q['coverage_at_eps'],q['avg_successors_at_eps'],q['smooth_sim']] for q in r],axis=0); print('beta',b,'mse %.5f hard %.5f coverage %.3f count %.3f sim %.5f'%tuple(vals))
121    print('math ratios',check['max_bound_ratio_min'],check['max_bound_ratio_max'],'errors',check['fixed_softmin_abs_error'])
122if __name__=='__main__': main()