<jupyter_start><jupyter_text>2d examples<jupyter_code>### Or choose cigar!!
from distributions import log_prob_banana as log_p

x = torch.linspace(-4, 4, 201)
xx, yy = torch.meshgrid(x, x)
xy = torch.stack([xx.flatten(), yy.flatten()], dim=0)
log_p_grid = log_p(xy)
p_grid = (log_p_grid - torch.logsumexp(log_p_grid.flatten(), dim=0)).exp().reshape(xx.size())
plt.contourf(xx, yy, p_grid, origin='lower', vmin=0., levels=50)<jupyter_output><empty_output><jupyter_text>Define a 2d polynomial<jupyter_code>n=10
random_coeffs = torch.randn(n,n)
f = torch.Tensor([[np.math.factorial(n) for n in range(10)]])
nicer_coeffs = random_coeffs/f/f.T

def random_poly(xy):
    """Basically np.polynomial.polynomial.polyval2d in pytorch.
    xy can be a (2,n) tensor"""
    if xy.size(0) == 2:
        x, y = xy[0], xy[1]
    else:
        x, y = xy[:, 0], xy[:, 1]
        
    terms = sum([[nicer_coeffs[i,j] * x**i * y**j for i in range(nicer_coeffs.size(0))]
                                                   for j in range(nicer_coeffs.size(1))], [])
    return torch.stack(terms, dim=0).sum(0) 
        
poly_grid = random_poly(torch.stack((xx,yy)))
plt.contourf(xx, yy, poly_grid, origin='lower', levels=50)
plt.colorbar()<jupyter_output><empty_output><jupyter_text>Get ground truth<jupyter_code>from scipy.integrate import dblquad
def integrand(x,y):
    xy = torch.Tensor([x, y])
    return random_poly(xy) * log_p(torch.Tensor(xy))
random_poly_expectation = dblquad(integrand, -10, 10, lambda x: -5, lambda x: 5)
random_poly_expectation<jupyter_output><empty_output><jupyter_text>Compute expectation under variational sol'n<jupyter_code>q_vi = MVNIso.new_random(d=2)
q_vi.theta.requires_grad_(True)

def _kl_q_p_helper():
    kl = -q_vi.entropy() - q_vi.quadrature_ev(log_p, 5)
    grad_kl = torch.autograd.grad(kl, q_vi.theta, create_graph=True)[0]
    hess = torch.zeros(q_vi.n_params, q_vi.n_params)
    for i in range(q_vi.n_params):
        hess[i, :] = torch.autograd.grad(grad_kl[i], q_vi.theta, retain_graph=True)[0]
    return kl.detach(), grad_kl.detach(), hess

kl_vals = torch.zeros(50)
for i in range(50):
    kl_vals[i], grad, hess = _kl_q_p_helper()
    with torch.no_grad():
        # Newton's method step
        q_vi.theta.copy_(q_vi.theta - torch.linalg.solve(hess, grad))
q_vi.theta.requires_grad_(False)

plt.figure()
plt.plot(kl_vals)
plt.xlabel('VI step')
plt.ylabel('KL(q||p) + log Z')
q_vi_expectation = q_vi.quadrature_ev(random_poly, 10)

plt.contourf(xx, yy, p_grid, origin='lower', vmin=0.)
plt.title('VI')
plt.xticks([]); 
plt.yticks([])
plt.axis('image')
plt.plot(*q_vi.ellipse(nsigma=1.), color='r', linewidth=1)
print(q_vi_expectation)<jupyter_output>tensor(1.3515)<jupyter_text>Do HMC sampling with various $\lambda$<jupyter_code>def log_psi(lam, q):
    kl_qp = -q.entropy() - q.quadrature_ev(log_p, n=25)
    return 0.5*q.log_det_fisher() - lam * kl_qp

def sample_theta(lam, n=1000):
    q = MVNIso(d=2)
    hmc = HMC(lambda th: log_psi(lam, q.set_theta(th)), leapfrog_t=2.)
    hmc.tune(torch.randn(200, q.n_params), target_accept=0.95)

    print(f"[lam={lam}]\tmass={hmc.mass.numpy()}\tdt={hmc.dt.item():.4f}")
    return hmc.sample(torch.randn(q.n_params), n, progbar=True)<jupyter_output><empty_output><jupyter_text>Run a long sampling chain once per $\lambda$<jupyter_code>#compute 
lambdas = torch.logspace(0,2,10)
hmc_results = []
for l in lambdas:
    save_file = Path('results') / f"{log_p.__name__}_hmc_lam={l:.3f}.dat"
    if save_file.exists():
        res = torch.load(save_file)
    else:
        res = sample_theta(l, n=10000)
        torch.save(res, save_file)
    hmc_results.append(res)
    print(f"\tfinished lam={l:.3f}, ESS={ess(res['samples']).numpy()}")
# or load from file
lambdas = torch.logspace(0,2,10)
hmc_results = []
for l in lambdas:
    save_file = Path('results') / f"{log_p.__name__}_hmc_lam={l:.3f}.dat"
    hmc_results.append(torch.load(save_file))<jupyter_output><empty_output><jupyter_text>Calculate 'asymptotic' bias by computing expected value random polynomial for every value of the chain<jupyter_code>q = MVNIso(d=2)
full_chain_expectations = torch.zeros(len(lambdas))
for i, r in enumerate(hmc_results):
    component_evs = [q.set_theta(th).quadrature_ev(random_poly, n=10) for th in r['samples']]
    full_chain_expectations[i] = torch.tensor(component_evs).mean()<jupyter_output><empty_output><jupyter_text>Calculate variance by subsampling each chain down to 100 samples many times<jupyter_code>n_subsample, n_run = 100, 50
q = MVNIso(d=2)
subsample_expectations = torch.zeros(len(lambdas), n_run)
for i, r in enumerate(hmc_results):
    for j in range(n_run):
        subset_samples = r['samples'][torch.randint(r['samples'].size(0), (n_subsample,)), :]
        component_evs = [q.set_theta(th).quadrature_ev(random_poly, n=10) for th in r['samples']]
        subsample_expectations[i, j] = torch.tensor(component_evs).mean()

est_variance = torch.mean((subsample_expectations - full_chain_expectations[:, None])**2)<jupyter_output><empty_output><jupyter_text>Plot EVs<jupyter_code>plt.semilogx(lambdas, full_chain_expectations, '-o')
plt.hlines(random_poly_expectation, 1,100,color='r')
plt.hlines(q_vi_expectation, 1,100,color='k')
plt.xlabel("$\lambda$");
plt.ylabel("EV")
plt.legend(['$m_\lambda(x)$', 'true value', 'VI EV')
plt.show()<jupyter_output><empty_output><jupyter_text>Plot bias$^2$ and variance<jupyter_code>bias2 = (full_chain_expectations - random_poly_expectation)**2
mse = bias2 + est_variance

plt.semilogx(lambdas, bias2, '-b', marker='.')
plt.semilogx(lambdas, est_variance, '-r', marker='.')
plt.semilogx(lambdas, mse, '-g', marker='.')
plt.xlabel("$\lambda$");
plt.legend(['bias$^2$', 'variance', 'MSE'])
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1660003476447071, 'fraction_numerical': 0.020337215365896055, 'mean_word_length': 5.247557003257329, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 40, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '363820', 'license': 'MIT', 'n_tokens_mistral': 2448, 'n_tokens_neox': 2342, 'n_words': 390, 'path': 'notebooks/bias_variance_hmc_2d.ipynb', 'repo_name': 'wrongu/sampling-variational', 'size': '5753'}