1import json, math, os
  2import numpy as np
  3
  4# Inverse-square reset diffusion toy experiment.
  5# Units: dt is the diffusion time step and D=1 unless specified.
  6def beta_pred(ratio):
  7    return (math.sqrt(1.0 + 4.0 * ratio) - 1.0) / 2.0
  8
  9def simulate(r0, lam=-2.0, D=1.0, eps=0.1, dt=0.002, nwalk=1800,
 10             steps=9000, burn=1500, seed=0, box=80.0):
 11    rng = np.random.default_rng(seed)
 12    z = np.zeros(nwalk, dtype=np.float64)
 13    samples = []
 14    reset_count = 0
 15    for t in range(steps):
 16        z += math.sqrt(2*D*dt) * rng.standard_normal(nwalk)
 17        # generalized hazard; for lambda=-2 this is exactly the proposal
 18        if lam == -2:
 19            rate = r0 / (z*z + eps*eps)
 20        else:
 21            rate = r0 / (np.abs(z) + eps) ** (-lam)
 22        q = 1.0 - np.exp(-dt * rate)
 23        hit = rng.random(nwalk) < q
 24        reset_count += int(hit.sum())
 25        z[hit] = 0.0
 26        # A large reflecting box avoids numerical overflow; box dependence is measured separately.
 27        z[z > box] = 2*box - z[z > box]
 28        z[z < -box] = -2*box - z[z < -box]
 29        if t >= burn and t % 10 == 0:
 30            samples.append(z.copy())
 31    x = np.abs(np.concatenate(samples))
 32    return x, reset_count / (nwalk * steps)
 33
 34def tail_fit(x, lo=2.0, hi=15.0):
 35    # CCDF ~ x^-(beta-1), so density exponent is one minus CCDF slope.
 36    xx = np.linspace(lo, hi, 40)
 37    cc = np.array([(x >= a).mean() for a in xx])
 38    good = (cc > 3e-4) & (cc < .3)
 39    if good.sum() < 5: return float('nan'), float('nan')
 40    slope, intercept = np.polyfit(np.log(xx[good]), np.log(cc[good]), 1)
 41    return 1.0 - slope, float(np.corrcoef(np.log(xx[good]), np.log(cc[good]))[0,1])
 42
 43def run_toy():
 44    ratios = [0.5, 1.0, 1.5, 2.0, 2.5, 4.0, 8.0]
 45    rows=[]
 46    for i,r in enumerate(ratios):
 47        x, rr = simulate(r, seed=100+i)
 48        b, fitq = tail_fit(x)
 49        rows.append({'r0_over_D':r, 'beta_pred':beta_pred(r), 'beta_fit':b,
 50                     'tail_fit_r':fitq, 'mean_abs':float(x.mean()),
 51                     'p_abs_gt_5':float((x>5).mean()), 'reset_rate_step':rr})
 52    # Boundary prediction: above ratio 2 has finite first moment only for beta>2,
 53    # while normalizability itself is beta>1. Compare probability mass in a growing box.
 54    box_rows=[]
 55    for r in [1.0, 4.0]:
 56        vals=[]
 57        for B in [10.,20.,40.]:
 58            x,_=simulate(r, box=B, steps=8000, burn=1200, seed=900+int(B)+int(r*10))
 59            vals.append({'box':B,'mean_abs':float(x.mean()),'p_gt_5':float((x>5).mean())})
 60        box_rows.append({'ratio':r,'beta_pred':beta_pred(r),'box_scaling':vals})
 61    # Controls for the claimed lambda classification, at r0=4.
 62    controls=[]
 63    for lam in [-3.,-2.,-1.]:
 64        x,_=simulate(4.,lam=lam,steps=9000,burn=1500,seed=700+int((lam+3)*10))
 65        b,q=tail_fit(x)
 66        controls.append({'lambda':lam,'mean_abs':float(x.mean()),'tail_density_exponent_fit':b,'fit_r':q})
 67    return {'toy_rows':rows,'box_rows':box_rows,'lambda_controls':controls}
 68
 69def run_ml(seed=42):
 70    # Small, reproducible regression; compare SGD with a checkpoint reset wrapper.
 71    import torch
 72    torch.manual_seed(seed); np.random.seed(seed)
 73    device='cuda' if torch.cuda.is_available() else 'cpu'
 74    try:
 75        dev=torch.device(device)
 76        x=torch.linspace(-3,3,512,device=dev).unsqueeze(1)
 77        y=torch.sin(3*x)+0.15*torch.randn_like(x)
 78        tr=x[:384]; ty=y[:384]; va=x[384:]; vy=y[384:]
 79        def train(reset=False):
 80            torch.manual_seed(seed)
 81            m=torch.nn.Sequential(torch.nn.Linear(1,32),torch.nn.Tanh(),torch.nn.Linear(32,1)).to(dev)
 82            opt=torch.optim.SGD(m.parameters(),lr=.035,momentum=.8)
 83            ref=[p.detach().clone() for p in m.parameters()]
 84            best=1e9; resets=0
 85            for step in range(700):
 86                idx=torch.randint(0,len(tr),(64,),device=dev)
 87                opt.zero_grad(); loss=((m(tr[idx])-ty[idx])**2).mean(); loss.backward(); opt.step()
 88                with torch.no_grad():
 89                    vl=((m(va)-vy)**2).mean().item()
 90                    if vl<best: best=vl; ref=[p.detach().clone() for p in m.parameters()]
 91                    if reset:
 92                        z=math.sqrt(sum(((p-r)**2).sum().item() for p,r in zip(m.parameters(),ref)))
 93                        # D estimate is deliberately fixed to the toy scale; this is a stress test, not tuning.
 94                        q=1-math.exp(-.004/(z*z+.05*.05))
 95                        if np.random.random()<q:
 96                            for p,r in zip(m.parameters(),ref): p.copy_(r)
 97                            opt.state.clear(); resets+=1
 98            return float(((m(va)-vy)**2).mean().item()),resets
 99        base=train(False); idea=train(True)
100        return {'device':str(dev),'baseline_val_mse':base[0],'reset_val_mse':idea[0],'reset_count':idea[1]}
101    except Exception as e:
102        return {'device':'cpu-fallback','error':repr(e)}
103
104if __name__=='__main__':
105    out={'toy':run_toy(),'ml':run_ml()}
106    with open('results.json','w') as f: json.dump(out,f,indent=2)
107    print(json.dumps(out,indent=2))