1import json, math
  2from pathlib import Path
  3import numpy as np
  4from scipy.stats import qmc, norm
  5from scipy.integrate import quad
  6
  7SEED = 2108
  8EPS = 1e-10
  9
 10# A one-step Euler probability-flow surrogate.  The monotone transport is
 11# T_lambda(z)=z+lambda*tanh(z), which has an exactly known change-of-variables
 12# Jacobian and is smooth enough for scrambled Sobol integration.
 13def transport(z, lam):
 14    return z + lam * np.tanh(z)
 15
 16def det_jacobian(z, lam):
 17    return np.prod(1.0 + lam / np.cosh(z)**2, axis=1)
 18
 19def log_phi(z):
 20    return -0.5*np.sum(z*z, axis=1) - z.shape[1]*0.5*math.log(2*math.pi)
 21
 22def log_pi(x):
 23    # Known target pi=N(0,I), so the correction is an exact importance ratio.
 24    return log_phi(x)
 25
 26def metric(x):
 27    return x[:, 0]**2
 28
 29def estimate(z, lam, weighted):
 30    x = transport(z, lam)
 31    vals = metric(x)
 32    if weighted:
 33        # pi(T(z))*|det JT(z)| / phi(z)
 34        logw = log_pi(x) + np.log(det_jacobian(z, lam)) - log_phi(z)
 35        vals = vals * np.exp(np.clip(logw, -40, 40))
 36    return float(np.mean(vals))
 37
 38def sobol_z(n, d, seed):
 39    u = qmc.Sobol(d=d, scramble=True, seed=seed).random_base2(int(round(math.log2(n))))
 40    return norm.ppf(np.clip(u, EPS, 1-EPS))
 41
 42def iid_z(n, d, rng):
 43    return rng.standard_normal((n, d))
 44
 45def exact_unweighted_expectation(lam):
 46    # E[(Z+lambda*tanh(Z))^2] = 1 + 2 lambda A + lambda^2 B.
 47    phi = lambda z: math.exp(-z*z/2)/math.sqrt(2*math.pi)
 48    A = quad(lambda z: z*math.tanh(z)*phi(z), -np.inf, np.inf, epsabs=1e-12)[0]
 49    B = quad(lambda z: math.tanh(z)**2*phi(z), -np.inf, np.inf, epsabs=1e-12)[0]
 50    return 1.0 + 2.0*lam*A + lam*lam*B, A, B
 51
 52def run():
 53    rng = np.random.default_rng(SEED)
 54    ns = [32, 64, 128, 256, 512, 1024]
 55    reps = 48
 56    z0, lam0, h = np.array([[0.37, -1.1]]), 0.8, 1e-5
 57    num = np.zeros((2, 2))
 58    for j in range(2):
 59        zp, zm = z0.copy(), z0.copy(); zp[0,j] += h; zm[0,j] -= h
 60        num[:,j] = ((transport(zp,lam0)-transport(zm,lam0))/(2*h))[0]
 61    jac_err = abs(np.linalg.det(num) - det_jacobian(z0,lam0)[0])
 62
 63    rows = []
 64    for lam in [0.0, 0.25, 0.5, 1.0]:
 65        exact_u, _, _ = exact_unweighted_expectation(lam)
 66        for n in ns:
 67            sob_u, sob_w, iid_u, iid_w = [], [], [], []
 68            for r in range(reps):
 69                z = sobol_z(n, 2, SEED + 10000*int(lam*100) + r)
 70                sob_u.append(estimate(z,lam,False)); sob_w.append(estimate(z,lam,True))
 71                zi = iid_z(n,2,rng)
 72                iid_u.append(estimate(zi,lam,False)); iid_w.append(estimate(zi,lam,True))
 73            rows.append(dict(lam=lam,n=n, exact_unweighted=exact_u,
 74                sobol_unweighted_mean=float(np.mean(sob_u)), sobol_unweighted_sd=float(np.std(sob_u,ddof=1)),
 75                sobol_weighted_mean=float(np.mean(sob_w)), sobol_weighted_sd=float(np.std(sob_w,ddof=1)),
 76                iid_unweighted_sd=float(np.std(iid_u,ddof=1)), iid_weighted_sd=float(np.std(iid_w,ddof=1)),
 77                sobol_unweighted_bias_to_target=float(np.mean(sob_u)-1),
 78                sobol_bias_to_exact_unweighted=float(np.mean(sob_u)-exact_u),
 79                sobol_weighted_bias=float(np.mean(sob_w)-1)))
 80
 81    def slope(lam, field):
 82        rr = [x for x in rows if x['lam']==lam]
 83        return float(np.polyfit(np.log2([x['n'] for x in rr]), np.log2([x[field] for x in rr]), 1)[0])
 84
 85    exact05, A, B = exact_unweighted_expectation(0.5)
 86    last = [x for x in rows if x['n']==1024]
 87    bias_fit = np.polyfit([x['lam'] for x in last], [x['sobol_unweighted_bias_to_target'] for x in last], 2)
 88    summary = {
 89      'jacobian_check_abs_error': float(jac_err),
 90      'target_expectation': 1.0,
 91      'analytic_unweighted_formula': '1 + 2*lambda*A + lambda^2*B',
 92      'analytic_A_E[z_tanh(z)]': A, 'analytic_B_E[tanh(z)^2]': B,
 93      'analytic_expectation_lambda_0.5': exact05,
 94      'slopes_sd_log2N_at_lambda_0.5': {
 95        'scrambled_sobol_unweighted': slope(0.5,'sobol_unweighted_sd'),
 96        'scrambled_sobol_weighted': slope(0.5,'sobol_weighted_sd'),
 97        'iid_unweighted': slope(0.5,'iid_unweighted_sd'),
 98        'iid_weighted': slope(0.5,'iid_weighted_sd')},
 99      'bias_quadratic_fit_coefficients_descending_lambda2_lambda_const': [float(v) for v in bias_fit],
100      'bias_vs_lambda_n1024': [dict(lam=x['lam'], predicted_bias=x['exact_unweighted']-1,
101          observed_bias=x['sobol_unweighted_bias_to_target'], weighted_bias=x['sobol_weighted_bias']) for x in last],
102      'rows': rows
103    }
104    Path('results.json').write_text(json.dumps(summary, indent=2))
105    print(json.dumps({k:v for k,v in summary.items() if k!='rows'}, indent=2))
106    print('TABLE lambda=0.5: n sobol_sd iid_sd weighted_sd target_bias weighted_bias')
107    for x in rows:
108        if x['lam']==0.5:
109            print(x['n'], '%.5g %.5g %.5g %.5g %.5g' % (x['sobol_unweighted_sd'],x['iid_unweighted_sd'],x['sobol_weighted_sd'],x['sobol_unweighted_bias_to_target'],x['sobol_weighted_bias']))
110
111if __name__ == '__main__':
112    run()