1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6from scipy.stats import qmc, norm
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report
  9
 10SEED=2108
 11TRACK='conditional_multitoken_diffusion'
 12MODEL='mlp_tiny'
 13LRS=[1e-3,3e-3,1e-2]
 14EPOCHS=12
 15BATCH=64
 16
 17def seed_all(s):
 18    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 19    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 20
 21def gaussian_corruption(n,d,seed,kind):
 22    if kind=='iid':
 23        z=np.random.RandomState(seed).standard_normal((n,d)).astype('float32')
 24    else:
 25        # Independent Owen scrambles per epoch/seed; inverse-CDF maps cube to N(0,1).
 26        p=qmc.Sobol(d=d, scramble=True, seed=int(seed)).random(n)
 27        z=norm.ppf(np.clip(p,1e-6,1-1e-6)).astype('float32')
 28    return torch.from_numpy(z)
 29
 30def run_one(seed,lr,kind,return_model=False):
 31    seed_all(seed)
 32    d=get_dataset(TRACK,seed,n_train=400,n_test=200)
 33    # The custom track exposes an 8-token target; adapt the harness' generic
 34    # regression reshape back to the declared multi-token output dimension.
 35    d['ytr']=d['ytr'].reshape(-1,8); d['yte']=d['yte'].reshape(-1,8)
 36    dev='cuda' if torch.cuda.is_available() else 'cpu'
 37    try:
 38        net=make_model(MODEL,d['input_shape'],d['out_dim']).to(dev)
 39        opt=torch.optim.Adam(net.parameters(),lr=lr)
 40        loss=nn.MSELoss()
 41        x=d['xtr'].to(dev); y=d['ytr'].to(dev)
 42        n=len(x)
 43        for ep in range(EPOCHS):
 44            # Same examples and targets; only diffusion perturbation differs.
 45            perm=torch.randperm(n,device=dev)
 46            xep=x[perm]; yep=y[perm]
 47            z=gaussian_corruption(n,8,seed*1000+ep,kind).to(dev)
 48            # Mild additional terminal-noise transport, with observed sigma retained.
 49            xin=xep.clone(); sig=xin[:,9:10].clamp(.04,1.)
 50            xin[:,:8]=xin[:,:8] + .18*sig*z
 51            for j in range(0,n,BATCH):
 52                opt.zero_grad(set_to_none=True)
 53                pred=net(xin[j:j+BATCH]); l=loss(pred,yep[j:j+BATCH])
 54                l.backward(); opt.step()
 55        with torch.no_grad():
 56            pred=net(d['xte'].to(dev)); metric=float(loss(pred,d['yte'].to(dev)).cpu())
 57        if return_model: return metric, net, d
 58        return metric
 59    except RuntimeError:
 60        # Explicit CPU fallback for constrained shared GPU slots.
 61        torch.cuda.empty_cache()
 62        osdev=torch.device('cpu')
 63        seed_all(seed)
 64        d=get_dataset(TRACK,seed,n_train=400,n_test=200)
 65        d['ytr']=d['ytr'].reshape(-1,8); d['yte']=d['yte'].reshape(-1,8)
 66        net=make_model(MODEL,d['input_shape'],d['out_dim']).to(osdev)
 67        opt=torch.optim.Adam(net.parameters(),lr=lr); loss=nn.MSELoss()
 68        x=d['xtr']; y=d['ytr']; n=len(x)
 69        for ep in range(EPOCHS):
 70            perm=torch.randperm(n); xin=x[perm].clone(); yep=y[perm]
 71            z=gaussian_corruption(n,8,seed*1000+ep,kind); xin[:,:8]+=0.18*xin[:,9:10].clamp(.04,1.)*z
 72            for j in range(0,n,BATCH):
 73                opt.zero_grad(); l=loss(net(xin[j:j+BATCH]),yep[j:j+BATCH]); l.backward(); opt.step()
 74        with torch.no_grad(): return float(loss(net(d['xte']),d['yte']))
 75
 76def fn(kind,lr): return lambda s: run_one(s,lr,kind)
 77
 78def mechanism_signature():
 79    # Re-test the stage-1 prediction on outputs of trained systems: batch-mean
 80    # variance should decay faster for scrambled Sobol than IID perturbations.
 81    seed=77; lr=3e-3; metric,net,d=run_one(seed,lr,'sobol',True)
 82    dev=next(net.parameters()).device; x=d['xte'][:128].to(dev)
 83    vals={}
 84    for kind in ('iid','sobol'):
 85        means=[]
 86        for r in range(16):
 87            xx=x.clone(); z=gaussian_corruption(len(x),8,90000+r,kind).to(dev)
 88            xx[:,:8]+=0.18*xx[:,9:10].clamp(.04,1.)*z
 89            with torch.no_grad(): means.append(float(net(xx).mean().cpu()))
 90        vals[kind]=float(np.std(means,ddof=1))
 91    # A second size check uses prefixes of power-of-two Sobol sets and IID sets.
 92    slopes={}
 93    for kind in ('iid','sobol'):
 94        ns=[16,32,64,128]; sds=[]
 95        for n in ns:
 96            mm=[]
 97            for r in range(12):
 98                xx=x[:n].clone(); z=gaussian_corruption(n,8,120000+r,kind).to(dev)
 99                xx[:,:8]+=0.18*xx[:,9:10].clamp(.04,1.)*z
100                with torch.no_grad(): mm.append(float(net(xx).mean().cpu()))
101            sds.append(np.std(mm,ddof=1))
102        slopes[kind]=float(np.polyfit(np.log2(ns),np.log2(np.maximum(sds,1e-12)),1)[0])
103    confirmed=vals['sobol'] < vals['iid'] and slopes['sobol'] < slopes['iid']
104    return {'prediction':'scrambled Sobol batch-mean variance lower and slope closer to -1 than IID',
105            'trained_model_batch_mean_sd':vals,'observed_log2N_slopes':slopes,
106            'variance_ratio_sobol_over_iid':vals['sobol']/vals['iid'], 'confirmed':bool(confirmed)}
107
108def main():
109    grid=[{'lr':v} for v in LRS]
110    base=sweep_baseline(lambda c: fn('iid',c['lr']),grid)
111    # Same union of learning rates; select best idea setting on the same 4-seed sweep.
112    idea_trials=[]
113    for cfg in grid:
114        r=evaluate(fn('sobol',cfg['lr']))
115        idea_trials.append({'cfg':cfg,'mean':r['mean'],'std':r['std']})
116    best=min(idea_trials,key=lambda x:x['mean'])['cfg']
117    idea=evaluate(fn('sobol',best['lr']))
118    rep=make_report(TRACK,MODEL,base,idea,{'idea_sweep':idea_trials,'selected_idea_cfg':best,
119                                           'mechanism_signature':mechanism_signature(),
120                                           'custom_track':{'name':TRACK,'file':'bench/custom_tracks/conditional_multitoken_diffusion.py','domain':'diffusion-sampling'}})
121    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
122    print(json.dumps(rep,indent=2))
123if __name__=='__main__': main()