1import json, math, random
  2import numpy as np
  3from scipy.special import ndtri
  4from scipy.stats import qmc
  5
  6SEED = 2437
  7np.random.seed(SEED); random.seed(SEED)
  8
  9
 10def sobol_base(n, d):
 11    return qmc.Sobol(d=d, scramble=False, seed=SEED).random_base2(int(round(math.log2(n))))
 12
 13
 14def qmc_estimates(n, d, reps, fn):
 15    x = sobol_base(n, d); out = []
 16    rng = np.random.default_rng(SEED + 17*d + n)
 17    for _ in range(reps):
 18        out.append(float(np.mean(fn((x + rng.random(d)) % 1.0))))
 19    return np.asarray(out)
 20
 21
 22def iid_estimates(n, d, reps, fn):
 23    rng = np.random.default_rng(SEED + 991*d + n)
 24    return np.asarray([float(np.mean(fn(rng.random((n,d))))) for _ in range(reps)])
 25
 26
 27def smooth(u):
 28    return np.prod(1.0 + 0.7*np.cos(2*np.pi*u), axis=1)
 29
 30
 31def nonsmooth(u):
 32    return (u[:,0] < 0.37).astype(float)
 33
 34
 35def toy_checks():
 36    unbiased = qmc_estimates(16, 2, 20000, smooth)
 37    ns = [4, 8, 16, 32, 64, 128, 256]
 38    smooth_rows=[]; rough_rows=[]
 39    for n in ns:
 40        q=qmc_estimates(n,2,512,smooth); i=iid_estimates(n,2,512,smooth)
 41        smooth_rows.append((n,float(np.sqrt(np.mean((q-1)**2))),float(np.sqrt(np.mean((i-1)**2)))))
 42        q=qmc_estimates(n,2,512,nonsmooth); i=iid_estimates(n,2,512,nonsmooth)
 43        rough_rows.append((n,float(np.sqrt(np.mean((q-.37)**2))),float(np.sqrt(np.mean((i-.37)**2)))))
 44    dim_rows=[]
 45    for d in [1,2,4,8,16]:
 46        q=qmc_estimates(64,d,512,smooth); i=iid_estimates(64,d,512,smooth)
 47        dim_rows.append((d,float(np.sqrt(np.mean((q-1)**2))),float(np.sqrt(np.mean((i-1)**2)))))
 48    def slope(rows,col):
 49        return float(np.polyfit(np.log([r[0] for r in rows]),np.log([max(r[col],1e-15) for r in rows]),1)[0])
 50    return {'unbiasedness':{'target':1.0,'mean':float(unbiased.mean()),'absolute_error':float(abs(unbiased.mean()-1)),'shift_sd':float(unbiased.std())},'smooth_scaling':{'rows_n_qmc_rmse_iid_rmse':smooth_rows,'qmc_loglog_slope':slope(smooth_rows,1),'iid_loglog_slope':slope(smooth_rows,2)},'nonsmooth_control':{'rows_n_qmc_rmse_iid_rmse':rough_rows,'qmc_loglog_slope':slope(rough_rows,1),'iid_loglog_slope':slope(rough_rows,2)},'dimension':{'rows_d_qmc_rmse_iid_rmse':dim_rows}}
 51
 52
 53def neural_gradient_check():
 54    import torch
 55    torch.manual_seed(SEED); dtype=torch.float64
 56    x0=torch.linspace(-1.,1.,48,dtype=dtype).reshape(-1,1); y=torch.sin(2.4*x0)+.15*x0
 57    net=torch.nn.Sequential(torch.nn.Linear(1,12),torch.nn.Tanh(),torch.nn.Linear(12,1)).double()
 58    state={k:v.detach().clone() for k,v in net.state_dict().items()}
 59    def grad_for(u):
 60        net.load_state_dict(state); net.zero_grad()
 61        z=torch.tensor(ndtri(np.clip(u,1e-6,1-1e-6)),dtype=dtype).reshape(-1,1)
 62        pred=net(x0)+.25*z.mean(); loss=((pred-y)**2).mean(); loss.backward()
 63        return np.concatenate([p.grad.detach().numpy().ravel() for p in net.parameters()]),float(loss.detach())
 64    n=16; reps=192; base=sobol_base(n,1); rng=np.random.default_rng(SEED+777)
 65    gq=[];gi=[];lq=[];li=[]
 66    for _ in range(reps):
 67        g,l=grad_for((base+rng.random(1))%1);gq.append(g);lq.append(l)
 68        g,l=grad_for(rng.random((n,1)));gi.append(g);li.append(l)
 69    gq=np.asarray(gq);gi=np.asarray(gi);qvar=float(np.mean(np.var(gq,axis=0,ddof=1)));ivar=float(np.mean(np.var(gi,axis=0,ddof=1)))
 70    return {'n':n,'reps':reps,'gradient_component_variance_qmc':qvar,'gradient_component_variance_iid':ivar,'variance_reduction':1-qvar/ivar,'loss_mean_qmc':float(np.mean(lq)),'loss_mean_iid':float(np.mean(li)),'loss_mean_difference':float(np.mean(lq)-np.mean(li))}
 71
 72
 73def training_check():
 74    import torch
 75    torch.manual_seed(SEED); dtype=torch.float64
 76    x=torch.linspace(-1.,1.,64,dtype=dtype).reshape(-1,1)
 77    target=lambda xx,zz: torch.sin(2.2*xx)+.25*zz
 78    base_model=torch.nn.Sequential(torch.nn.Linear(2,16),torch.nn.Tanh(),torch.nn.Linear(16,1)).double()
 79    init={k:v.detach().clone() for k,v in base_model.state_dict().items()}
 80    n,steps=16,120; sob=sobol_base(n,2); rng=np.random.default_rng(SEED+314)
 81    curves=[]
 82    for mode in ['qmc','iid']:
 83        model=torch.nn.Sequential(torch.nn.Linear(2,16),torch.nn.Tanh(),torch.nn.Linear(16,1)).double(); model.load_state_dict(init)
 84        opt=torch.optim.Adam(model.parameters(),lr=.025); vals=[]
 85        for step in range(steps):
 86            if mode=='qmc': u=(sob+rng.random(2))%1
 87            else: u=rng.random((n,2))
 88            z=torch.tensor(ndtri(np.clip(u[:,1],1e-6,1-1e-6)),dtype=dtype).reshape(-1,1)
 89            # Use the first coordinate to select fixed data points, while the second is latent noise.
 90            idx=np.floor(u[:,0]*len(x)).astype(int); xx=x[idx]; yy=target(xx,z)
 91            pred=model(torch.cat([xx,z],1)); loss=((pred-yy)**2).mean()
 92            opt.zero_grad();loss.backward();opt.step();vals.append(float(loss.detach()))
 93        curves.append(vals)
 94    return {'steps':steps,'final_loss_qmc':curves[0][-1],'final_loss_iid':curves[1][-1],'best_loss_qmc':min(curves[0]),'best_loss_iid':min(curves[1]),'trajectory_first_last':{'qmc':[curves[0][0],curves[0][-1]],'iid':[curves[1][0],curves[1][-1]]}}
 95
 96
 97def main():
 98    result={'seed':SEED,'toy':toy_checks(),'neural_gradient':neural_gradient_check(),'training':training_check()}
 99    with open('results.json','w') as f: json.dump(result,f,indent=2)
100    print(json.dumps(result,indent=2))
101
102if __name__=='__main__': main()