<jupyter_start><jupyter_text>Mueller potential---<jupyter_code>#paper_dir = '/Users/noe/data/papers/BoltzmannGenerators/'
params = MuellerPotential.params_default.copy()
params['k'] = 0.1
params['dim'] = 2
mueller = MuellerPotential(params=params)
#plt.figure(figsize=(5,5))
#double_well.plot_dimer_energy();
#plt.savefig(paper_dir + 'figs/doublewell_potential.pdf', bbox_inches='tight')
def plot_potential(labels=True, cbar=True, figsize=(7, 4)):
    # 2D potential
    xgrid = np.linspace(-1.5, 1.0, 100)
    ygrid = np.linspace(-0.2, 2, 100)
    Xgrid, Ygrid = np.meshgrid(xgrid, ygrid)
    X = np.vstack([Xgrid.flatten(), Ygrid.flatten(), np.zeros(Xgrid.size), np.zeros(Xgrid.size)]).T
    E = mueller.energy(X)
    E = E.reshape((100, 100))
    E = np.minimum(E, 10.0)
    plt.figure(figsize=figsize)
    plt.contourf(Xgrid, Ygrid, E, 100, cmap='jet', vmin=-10, vmax=-3)
    if cbar:
        cbar = plt.colorbar()
        cbar.set_label('Energy / kT') # , labelpad=-15, y=0.6
        #cbar.set_ticks([-10, -5, 0, 5, 10])
    if labels:
        plt.xlabel('$x_1$ / a.u.')
        plt.ylabel('$x_2$ / a.u.')
    else:
        plt.xticks([])
        plt.yticks([])
plot_potential()
# simulation data
from deep_boltzmann.sampling import MetropolisGauss
nsteps = 10000
x0_left = np.array([[-0.5, 1.5]])
x0_right = np.array([[0.5, 0.0]])

sampler = MetropolisGauss(mueller, x0_left, noise=0.02, stride=10)
sampler.run(nsteps)
traj_left = sampler.traj.copy()

sampler.reset(x0_left)
sampler.run(nsteps)
traj_left_val = sampler.traj.copy()

sampler.reset(x0_right)
sampler.run(nsteps)
traj_right = sampler.traj.copy()

sampler.reset(x0_right)
sampler.run(nsteps)
traj_right_val = sampler.traj.copy()
plot_potential(cbar=False, labels=False)
plt.plot(traj_left[:, 0], traj_left[:, 1])
plt.plot(traj_right[:, 0], traj_right[:, 1])
#plt.plot(traj_left_val[:, 0], traj_left_val[:, 1])
#plt.plot(traj_right_val[:, 0], traj_right_val[:, 1])
x = np.vstack([traj_left, traj_right])
xval = np.vstack([traj_left_val, traj_right_val])
plt.plot(x[:, 0], x[:, 1], linewidth=0, marker='.')
plt.plot(xval[:, 0], xval[:, 1], linewidth=0, marker='.')
#plt.plot(traj_right[:, 0], traj_right[:, 1])<jupyter_output><empty_output><jupyter_text>Particle filter starting from one sample-----<jupyter_code>x0 = traj_left[100:101]
X0 = np.repeat(x0, 1000, axis=0)
#X0 += 0.01 * np.random.randn(X0.shape[0], X0.shape[1])
plot_potential(labels=False, cbar=False, figsize=(5, 5))
plt.plot(x0[:, 0], x0[:, 1], color='white', linewidth=0, marker='+', markersize=20, markeredgewidth=3)
#plt.savefig(paper_dir + 'figs/mueller/explore_potential_init.pdf', bbox_inches='tight', transparent=True)
bg = invnet(mueller.dim, 'RRRRR', mueller, nl_layers=3, nl_hidden=100, # RPRPRPRPR
            nl_activation='relu', nl_activation_scale='tanh')
loss_bg_trainML, loss_bg_valML = bg.train_ML(X0, epochs=3, batch_size=128, 
                                             std=1.0, verbose=0, return_test_energies=False)
from deep_boltzmann.networks.training import ParticleFilter
particle_filter = ParticleFilter(bg, X0, 100000, lr=0.0001, batch_size=2048, 
                                 high_energy=10000, max_energy=1e10, std=1.0, weigh_ML=True)
Ds = []
Ds.append(proj_diag(particle_filter.X))
for i in range(30):
    print('\nITER',(i+1))
    particle_filter.train(epochs=50, stepsize=None, verbose=1)
    Ds.append(proj_diag(particle_filter.X))
fig, axes = plot_convergence(particle_filter.loss_train, particle_filter.acceptance_rate, particle_filter.stepsize[1:], figsize=(5, 5));
axes[0].set_ylim(-5, 5)
axes[1].set_ylim(-10, 40)
#plt.savefig(paper_dir + 'figs/mueller/explore_conv.pdf', bbox_inches='tight')
#particle_filter.stepsize = list(np.concatenate([np.repeat(0.1, 61), particle_filter.stepsize]))
len(particle_filter.stepsize)
plot_indices = np.array([0, 1, 2, 5, 20, 30])
x_true, p_true = proj_equilibrium_density()
fig, axes = plt.subplots(nrows=3, ncols=2, sharex=True, sharey=True, figsize=(6,8))
for i, ax in zip(plot_indices, axes.flatten()):
    h, e = np.histogram(Ds[i], bins=50)
    e = 0.5*(e[:-1] + e[1:])
    ax.plot(e, h, color='darkblue', linewidth=1)
    ax.fill_between(e, np.zeros(len(e)), h, color='blue', alpha=0.7)
    ax.text(-1, 20000, str(i*50)+' iter')
    ax.plot(x_true, 10*p_true, color='black', linewidth=3, label='Eq. dist.')
    ax.semilogy()
    #plt.legend(ncol=1, fontsize=16, frameon=False)
    ax.set_ylim(3, 100000)
axes[-1, 0].set_xlabel('$x_1$')
axes[-1, 1].set_xlabel('$x_1$')
axes[0, 0].set_ylabel('sample count')
axes[1, 0].set_ylabel('sample count')
axes[2, 0].set_ylabel('sample count')
#plt.savefig(paper_dir + 'figs/mueller/explore_conv_hist.pdf', bbox_inches='tight')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18472958642629905, 'fraction_numerical': 0.0496288441145281, 'mean_word_length': 4.977186311787072, 'pattern_counts': {'":': 0, '<': 9, '<?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': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '242492', 'license': 'CC-BY-4.0', 'n_tokens_mistral': 2118, 'n_tokens_neox': 1948, 'n_words': 316, 'path': 'notebooks/Fig4_Explore_Mueller.ipynb', 'repo_name': 'TrentBrick/learning-protein-landscapes', 'size': '4715'}