<jupyter_start><jupyter_text>Pick a few values of $\lambda$ and, for each one, configure the bounds of $\theta$-space that we will discretize<jupyter_code>configs = [{'lambda': 1.0, 'log_sig_range': (-32, 8), 'mu_range': (-5, 5)},
           {'lambda': 2.0, 'log_sig_range': (-2.5, 1.5), 'mu_range': (-5, 5)},
           {'lambda': 8.0, 'log_sig_range': (-2.5, 1.5), 'mu_range': (-5, 5)},
           {'lambda': 64.0, 'log_sig_range': (-2.5, 1.5), 'mu_range': (-5, 5)}]
q = MVNIso(d=1)
for conf in configs:
    conf['gam_grid'] = torch.linspace(*conf['log_sig_range'], 200)
    mm, gg = torch.meshgrid(x_grid, conf['gam_grid'])
    theta_grid = torch.stack([mm.flatten(), gg.flatten()], dim=1)
    
    conf['logdet_fisher_grid'] = torch.tensor([q.set_theta(th).log_det_fisher() for th in theta_grid])

    conf['ent_q_grid'] = torch.tensor([q.set_theta(th).entropy() for th in theta_grid])
    conf['ce_qp_grid'] = torch.tensor([-q.set_theta(th).quadrature_ev(log_p, 25) for th in theta_grid])
    conf['kl_qp_grid'] = -conf['ent_q_grid'] + conf['ce_qp_grid']
    
    conf['discrete_log_psi'] = 0.5*conf['logdet_fisher_grid'] - conf['lambda'] * conf['kl_qp_grid']<jupyter_output><empty_output><jupyter_text>Figure 2<jupyter_code>def sample_theta(psi, n, mm, gg):
    idx = torch.multinomial(psi.flatten(), num_samples=n, replacement=True)
    return torch.stack([mm.flatten()[idx], gg.flatten()[idx]], dim=1)

def log2prob(logp):
    p = torch.exp(logp - logp.max())
    return p / p.sum()

def log_q(x, theta):
    sig2 = torch.exp(2*theta[1])
    return -0.5*(x.view(-1,1)-theta[0])**2/sig2 - log(2*pi)/2 - theta[1]
n_example_components = 10
fig, ax = plt.subplots(len(configs), 2, figsize=(6, 3*len(configs)) )
for i, c in enumerate(configs):
    psi = log2prob(c['discrete_log_psi'])
    mm, gg = torch.meshgrid(x_grid, c['gam_grid'])
    th_samples = sample_theta(psi, n_example_components, mm, gg)
    log_q_max = -c['gam_grid'].min()  # log of peak height of the narrowest q kernel is roughly -log(sigma) for the min sigma
    discrete_m = sum((log_q(x_grid, (m, g)) - log_q_max).exp()*p for m, g, p in zip(mm.flatten(), gg.flatten(), psi.flatten()))
    discrete_m = discrete_m / discrete_m.sum()
    
    # PLOT1 : p, q, m
    log_q_values = log_q(x_grid, th_samples.T)
    q_values = log_q_values.exp() / log_q_values.exp().sum(dim=0, keepdims=True)
    ax[i,0].plot(x_grid, discrete_m, '-b')
    ax[i,0].plot(x_grid, discrete_p, '--k')

    yl = ax[i,0].get_ylim()
    for j in range(n_example_components):
        ax[i,0].fill_between(x_grid, q_values[:,j], color=(1.,0.,0.,.1))
    ax[i,0].set_yticks([])
    ax[i,0].set_ylim(yl)
    ax[i,0].set_xlabel('x', fontsize=15)
    ax[i,0].set_title("$\lambda=$"+str(int(c['lambda'])), fontsize=18)
    # PLOT2 : psi
#     print(psi.max(), psi.median())
    ax[i,1].contourf(mm, gg/log(10), psi.reshape(mm.size()), origin='lower', vmin=0.,)
    ax[i,1].plot(th_samples[:,0], th_samples[:,1]/log(10), '.r')
    ax[i,1].set_xlabel(r'$\mu$', fontsize=15)
    ax[i,1].set_ylabel(r'$\log_{10}(\sigma)$', fontsize=15)
    ax[i,1].set_xlim(c['mu_range'])
    if i==0:
        ax[i,1].annotate("$-\infty$", (.1,.10), (.08,.15), arrowprops={'arrowstyle':'->','color':'w'},
                         xycoords ='subfigure fraction', c='w', fontsize=15)
fig.tight_layout()
fig.savefig('mi_kl_figure_1d.svg')<jupyter_output>/var/folders/ph/lx9dyyx91893rxnlw09n7t0c0000gn/T/ipykernel_20063/332499431.py:34: UserWarning: Tight layout not applied. tight_layout cannot make axes height small enough to accommodate all axes decorations
  fig.tight_layout()
{'fraction_non_alphanumeric': 0.20478575403450194, 'fraction_numerical': 0.03951029493600445, 'mean_word_length': 4.255847953216374, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '661583', 'license': 'MIT', 'n_tokens_mistral': 1559, 'n_tokens_neox': 1488, 'n_words': 263, 'path': 'notebooks/mi_kl_figure_1d.ipynb', 'repo_name': 'wrongu/sampling-variational', 'size': '3594'}