1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 3133
  6rng = np.random.default_rng(SEED)
  7D = 8
  8N = 64
  9K = 120
 10
 11# Eight separated wells in an 8-D neural-energy-like landscape.
 12centers = np.zeros((8, D))
 13for j in range(8):
 14    centers[j, j] = 2.4
 15    centers[j, (j + 1) % D] = -1.4
 16sigma = 0.72
 17log2pi = math.log(2 * math.pi)
 18
 19def logp0(x):
 20    return -0.5 * np.sum(x*x, axis=-1) - D*0.5*log2pi
 21
 22def score(x):
 23    # log likelihood of a mixture of narrow wells; max/logsumexp is stable.
 24    z = -0.5*np.sum((x[:, None, :] - centers[None, :, :])**2, axis=-1)/sigma**2
 25    z -= D*math.log(sigma) + D*0.5*log2pi + math.log(8.)
 26    m = z.max(axis=1)
 27    return m + np.log(np.exp(z-m[:,None]).sum(axis=1))
 28
 29def log_gaussian(x, mean, cov):
 30    L = np.linalg.cholesky(cov)
 31    y = np.linalg.solve(L, (x-mean).T).T
 32    return -0.5*np.sum(y*y, axis=1) - np.log(np.diag(L)).sum() - D*0.5*log2pi
 33
 34def systematic_resample(w, rr):
 35    w = w / w.sum()
 36    c = np.cumsum(w)
 37    return int(np.searchsorted(c, rr))
 38
 39def importance_sanity():
 40    # Deliberately biased proposal. Weighted moments should recover constrained-prior moments.
 41    M = 180000
 42    qmean = np.ones(D) * 1.3
 43    qcov = np.eye(D) * 2.0
 44    x = rng.multivariate_normal(qmean, qcov, size=M)
 45    s = score(x)
 46    lam = np.quantile(s, .88)
 47    lpq = log_gaussian(x, qmean, qcov)
 48    lp = logp0(x)
 49    ok = s > lam
 50    lw = lp[ok] - lpq[ok]
 51    lw -= lw.max()
 52    w = np.exp(lw)
 53    ess = w.sum()**2 / np.sum(w*w)
 54    # Independent prior sample gives a reference constrained mean (not used in weighting).
 55    y = rng.normal(size=(500000, D)); sy = score(y); yy = y[sy > lam]
 56    weighted_mean = (w[:,None]*x[ok]).sum(0)/w.sum()
 57    ref_mean = yy.mean(0)
 58    mean_err = float(np.linalg.norm(weighted_mean-ref_mean)/math.sqrt(D))
 59    return dict(threshold=float(lam), ess=float(ess), ess_fraction=float(ess/max(1,ok.sum())),
 60                weighted_reference_rmse=mean_err, accepted=int(ok.sum()))
 61
 62def beta_sanity():
 63    # Exact claimed shrinkage law, with a sizeable independent replication.
 64    u = rng.random((200000, N))
 65    logs = np.log(u.max(axis=1))
 66    return dict(empirical_mean=float(logs.mean()), predicted=-1/N,
 67                relative_error=float(abs(logs.mean()+1/N)/(1/N)),
 68                empirical_sd=float(logs.std()), predicted_sd=float(1/N))
 69
 70def fit_flow(live):
 71    # Conditional affine flow: ML Gaussian location/scale fitted on recent live set.
 72    # Shrinkage/regularization prevents singular densities in a tiny online window.
 73    mu = live.mean(0)
 74    var = live.var(0) + 0.20**2
 75    return mu, var
 76
 77def run_nested(method):
 78    local = np.random.default_rng(SEED + (0 if method=='baseline' else 1))
 79    live = local.normal(size=(N,D)); ls = score(live)
 80    evals = 0; accepted = 0; ess_values=[]; acc_values=[]; trajectories=[]
 81    logX = 0.0
 82    for it in range(K):
 83        worst = int(np.argmin(ls)); lam = ls[worst]
 84        # For baseline, exact prior rejection. For idea, batch proposals from fitted affine flow.
 85        if method == 'baseline':
 86            ntry = 0
 87            while True:
 88                x = local.normal(size=(1,D)); sx = score(x)[0]; ntry += 1
 89                if sx > lam: break
 90            evals += ntry; acc_values.append(1.0/ntry); ess_values.append(1.0)
 91        else:
 92            mu, var = fit_flow(live)
 93            # A modest batch keeps the experiment small while allowing importance resampling.
 94            M = 128
 95            x = mu + local.normal(size=(M,D))*np.sqrt(var)
 96            sx = score(x); evals += M
 97            ok = sx > lam
 98            if not np.any(ok):
 99                # Robust fallback: exact prior rejection for this replacement.
100                ntry=0
101                while True:
102                    xx=local.normal(size=(1,D)); ss=score(xx)[0]; ntry+=1
103                    if ss>lam: x=xx; sx=np.array([ss]); break
104                evals += ntry; accepted += 1; live[worst]=x[0]; ls[worst]=sx[0]
105                ess_values.append(0.0); acc_values.append(0.0); logX += -1/N; trajectories.append(logX); continue
106            lpq = log_gaussian(x[ok], mu, np.diag(var))
107            lp = logp0(x[ok])
108            lw = lp-lpq; lw -= lw.max(); w=np.exp(lw)
109            ess_values.append(float(w.sum()**2/np.sum(w*w)))
110            acc_values.append(float(ok.mean()))
111            j = systematic_resample(w, local.random())
112            xnew=x[ok][j]; snew=sx[ok][j]
113        live[worst]=xnew if method!='baseline' else x[0]
114        ls[worst]=snew if method!='baseline' else sx
115        accepted += 1
116        logX += -1/N # ideal nested trajectory comparator
117        trajectories.append(logX)
118    return dict(score_evals=evals, accepted=accepted,
119                evals_per_accept=evals/accepted, mean_candidate_acceptance=float(np.mean(acc_values)),
120                mean_ess=float(np.mean(ess_values)), final_logX=float(logX),
121                trajectory=trajectories)
122
123def main():
124    t=time.time(); out={'seed':SEED, 'dimension':D, 'live_points':N, 'iterations':K,
125        'math_checks':{'beta_shrinkage':beta_sanity(), 'importance_sampling':importance_sanity()}}
126    out['baseline']=run_nested('baseline'); out['idea']=run_nested('idea')
127    # Avoid bloating the report while retaining a trajectory snapshot for inspection.
128    out['baseline']['trajectory_sample']=out['baseline'].pop('trajectory')[::20]
129    out['idea']['trajectory_sample']=out['idea'].pop('trajectory')[::20]
130    out['runtime_sec']=time.time()-t
131    Path('results.json').write_text(json.dumps(out, indent=2))
132    print(json.dumps(out, indent=2))
133
134if __name__=='__main__': main()