<jupyter_start><jupyter_text>Testing Score-Only MCMC sampling<jupyter_code>%pylab inline
from nsec.samplers import ScoreHamiltonianMonteCarlo, ScoreMetropolisAdjustedLangevinAlgorithm
from nsec.datasets.swiss_roll import get_swiss_roll
import jax
import jax.numpy as jnp
import tensorflow_probability as tfp; tfp = tfp.experimental.substrates.jax<jupyter_output>Populating the interactive namespace from numpy and matplotlib<jupyter_text>Instantiating test distribution<jupyter_code># Create a swiss roll distribution
dist = get_swiss_roll(.5, 2048)

# Define the score function
score = jax.grad(dist.log_prob)
init_samples = dist.sample(10000, seed=jax.random.PRNGKey(0))
hist2d(init_samples[:,0], init_samples[:,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')<jupyter_output><empty_output><jupyter_text>Validating Score-Hamiltonian Monte-Carlo<jupyter_code># Initialize the HMC transition kernel.
num_results = int(20e3)
num_burnin_steps = int(1e3)

# First running SHMC
kernel_shmc = ScoreHamiltonianMonteCarlo(
        target_score_fn=jax.grad(dist.log_prob),
        num_leapfrog_steps=10,
        num_delta_logp_steps=16,
        step_size=0.1)

samples_shmc, is_accepted_shmc = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=init_samples[1],
      kernel=kernel_shmc,
      trace_fn=lambda _, pkr: pkr.is_accepted, 
      seed=jax.random.PRNGKey(1))
# Then running true HMC
kernel_hmc = ScoreHamiltonianMonteCarlo(
        target_score_fn=jax.grad(dist.log_prob),
        num_leapfrog_steps=10,
        num_delta_logp_steps=16,
        step_size=0.1)

samples_hmc, is_accepted_hmc = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=init_samples[1],
      kernel=kernel_hmc,
      trace_fn=lambda _, pkr: pkr.is_accepted, 
      seed=jax.random.PRNGKey(1))
figure(figsize=(15,5))
subplot(131)
hist2d(init_samples[:,0], init_samples[:,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Samples from true distribution')
legend()
subplot(132)
hist2d(samples_shmc[is_accepted_shmc,0], samples_shmc[is_accepted_shmc,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Score-only HMC')
legend()
subplot(133)
hist2d(samples_hmc[is_accepted_hmc,0], samples_hmc[is_accepted_hmc,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('True HMC')
legend()
# No difference in number of accepted steps
print(len(where(is_accepted_hmc)[0]), len(where(is_accepted_shmc)[0]))
# Reducing the number of points in the logp integration
kernel_shmc = ScoreHamiltonianMonteCarlo(
        target_score_fn=jax.grad(dist.log_prob),
        num_leapfrog_steps=10,
        num_delta_logp_steps=2,
        step_size=0.1)

samples_shmc, is_accepted_shmc = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=init_samples[1],
      kernel=kernel_shmc,
      trace_fn=lambda _, pkr: pkr.is_accepted, 
      seed=jax.random.PRNGKey(1))
figure(figsize=(15,5))
subplot(131)
hist2d(init_samples[:,0], init_samples[:,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Samples from true distribution')
legend()
subplot(132)
hist2d(samples_shmc[is_accepted_shmc,0], samples_shmc[is_accepted_shmc,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Score-only HMC (2 pts integration)')
legend()
subplot(133)
hist2d(samples_hmc[is_accepted_hmc,0], samples_hmc[is_accepted_hmc,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('True HMC')
legend()
# Even with 2 integration points for delta logp, no obvious problem
print(len(where(is_accepted_hmc)[0]), len(where(is_accepted_shmc)[0]))<jupyter_output>19947 19880<jupyter_text>Validating Score MALA<jupyter_code># First running SMALA
kernel_smala = ScoreMetropolisAdjustedLangevinAlgorithm(
        target_score_fn=jax.grad(dist.log_prob),
        num_delta_logp_steps=16,
        step_size=0.1)

samples_smala, is_accepted_smala = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=init_samples[1],
      kernel=kernel_smala,
      trace_fn=lambda _, pkr: pkr.is_accepted, 
      seed=jax.random.PRNGKey(1))
kernel_mala = tfp.mcmc.MetropolisAdjustedLangevinAlgorithm(
        target_log_prob_fn=dist.log_prob,
        step_size=0.1)

samples_mala, is_accepted_mala = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=init_samples[1],
      kernel=kernel_mala,
      trace_fn=lambda _, pkr: pkr.is_accepted, 
      seed=jax.random.PRNGKey(1))
figure(figsize=(15,5))
subplot(131)
hist2d(init_samples[:,0], init_samples[:,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Samples from true distribution')
legend()
subplot(132)
hist2d(samples_smala[is_accepted_smala,0], samples_smala[is_accepted_smala,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Score-only MALA (16 pts integration)')
legend()
subplot(133)
hist2d(samples_mala[is_accepted_mala,0], samples_mala[is_accepted_mala,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('True MALA')
legend()
print(len(where(is_accepted_mala)[0]), len(where(is_accepted_smala)[0]))
# Testing with fewer integration steps
kernel_smala = ScoreMetropolisAdjustedLangevinAlgorithm(
        target_score_fn=jax.grad(dist.log_prob),
        num_delta_logp_steps=2,
        step_size=0.1)

samples_smala, is_accepted_smala = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=init_samples[1],
      kernel=kernel_smala,
      trace_fn=lambda _, pkr: pkr.is_accepted, 
      seed=jax.random.PRNGKey(1))
figure(figsize=(15,5))
subplot(131)
hist2d(init_samples[:,0], init_samples[:,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Samples from true distribution')
legend()
subplot(132)
hist2d(samples_smala[is_accepted_smala,0], samples_smala[is_accepted_smala,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('Score-only MALA (2 pts integration)')
legend()
subplot(133)
hist2d(samples_mala[is_accepted_mala,0], samples_mala[is_accepted_mala,1],64,range=[[-15,15],[-15,15]]); gca().set_aspect('equal')
scatter(init_samples[1,0], init_samples[1,1], label='x0', color='C1')
title('True MALA')
legend()
print(len(where(is_accepted_mala)[0]), len(where(is_accepted_smala)[0]))<jupyter_output>19555 19574
{'fraction_non_alphanumeric': 0.15407813570939, 'fraction_numerical': 0.051542152159013026, 'mean_word_length': 7.097669256381798, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 14, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '492468', 'license': 'MIT', 'n_tokens_mistral': 3008, 'n_tokens_neox': 2863, 'n_words': 371, 'path': 'notebooks/test_samplers.ipynb', 'repo_name': 'b-remy/score-estimation-comparison', 'size': '7295'}