def jump_plot(distribution, nsamples=100, **kwargs):
"""
Plots samples drawn from distribution with dwelling time on the x-axis
and the sample value on the y-axis
1D only
"""
distr = distribution(ndims=1, nbatch=1, **kwargs)
# sampler = MarkovJumpHMC(np.array([0]).reshape(1,1),
sampler = MarkovJumpHMC(distr.Xinit, distr.E, distr.dEdX, epsilon=0.3, beta=0.2, num_leapfrog_steps=5)
x_t = []
d_t = []
transitions = []
last_L_count, last_F_count, last_R_count = 0, 0, 0
for idx in xrange(nsamples):
sampler.sampling_iteration()
x_t.append(sampler.state.X[0, 0])
d_t.append(sampler.dwelling_times[0])
if sampler.L_count - last_L_count == 1:
transitions.append("L")
last_L_count += 1
elif sampler.F_count - last_F_count == 1:
transitions.append("F")
last_F_count += 1
elif sampler.R_count - last_R_count == 1:
transitions.append("R")
last_R_count += 1
t = np.cumsum(d_t)
plt.scatter(t, x_t)
t = np.array(t).reshape(len(t), 1)
x_t = np.array(x_t).reshape(len(x_t), 1)
transitions = np.array(transitions).reshape(len(transitions), 1)
data = np.concatenate((x_t, t, transitions), axis=1)
return pd.DataFrame(data, columns=["x", "t", "transitions"])
def generate_initialization(distribution):
""" Run mjhmc for BURN_IN_STEPS on distribution, generating a fair set of initial states
:param distribution: Distribution object. Must have nbatch == MAX_N_PARTICLES
:returns: a set of fair initial states and an estimate of the variance for emc and true both
:rtype: tuple: (array of shape (distribution.ndims, MAX_N_PARTICLES), float, float)
"""
print('Generating fair initialization for {} by burning in {} steps'.format(
type(distribution).__name__, BURN_IN_STEPS))
assert BURN_IN_STEPS > VAR_STEPS
assert distribution.nbatch == MAX_N_PARTICLES
mjhmc = MarkovJumpHMC(distribution=distribution, resample=False)
for _ in xrange(BURN_IN_STEPS - VAR_STEPS):
mjhmc.sampling_iteration()
assert mjhmc.resample == False
emc_var_estimate, mjhmc = online_variance(mjhmc, distribution)
# we discard v since p(x,v) = p(x)p(v)
fair_x = mjhmc.state.copy().X
# otherwise will go into recursive loop
distribution.mjhmc = False
control = ControlHMC(distribution=distribution.reset())
for _ in xrange(BURN_IN_STEPS - VAR_STEPS):
control.sampling_iteration()
true_var_estimate, control = online_variance(control, distribution)
return (fair_x, emc_var_estimate, true_var_estimate)
def sampler_speedometer():
""" Tests average sample speed of different samplers in different configurations
:returns: Just prints info
:rtype: None
"""
gaussian = Gaussian()
np.random.seed(2015)
pot = ProductOfT(ndims=36,nbasis=36)
mjhmc_gauss = MarkovJumpHMC(distribution=gaussian)
mjhmc_gauss_nr = MarkovJumpHMC(distribution=gaussian, resample=False)
control_gauss = ControlHMC(distribution=gaussian)
mjhmc_pot = MarkovJumpHMC(distribution=pot)
mjhmc_pot_nr = MarkovJumpHMC(distribution=pot, resample=False)
control_pot = ControlHMC(distribution=pot)
m_g_r_avg = time_per_sample(mjhmc_gauss)
m_g_nr_avg = time_per_sample(mjhmc_gauss_nr)
c_g_avg = time_per_sample(control_gauss)
m_p_r_avg = time_per_sample(mjhmc_pot)
m_p_nr_avg = time_per_sample(mjhmc_pot_nr)
c_p_avg = time_per_sample(control_pot)
print "Average times per samples..."
print "resampled MJHMC numpy gradient: {}".format(m_g_r_avg)
print "not resampled MJHMC numpy gradient: {}".format(m_g_nr_avg)
print "control HMC numpy gradient: {}".format(c_g_avg)
print "resampled MJHMC theano gradient: {}".format(m_p_r_avg)
print "not resampled MJHMC theano gradient: {}".format(m_p_nr_avg)
print "control HMC theano gradient: {}".format(c_p_avg)
def hist_2d(distribution, nsamples, **kwargs):
"""
Plots a 2d hexbinned histogram of distribution
"""
distr = distribution(ndims=2)
sampler = MarkovJumpHMC(distr.Xinit, distr.E, distr.dEdX, **kwargs)
samples = sampler.sample(nsamples)
with sns.axes_style("white"):
sns.jointplot(samples[0], samples[1], kind="kde", stat_func=None)
def hist_2d(distr, nsamples, **kwargs):
"""
Plots a 2d hexbinned histogram of distribution
Args:
distr: Distribution object
nsamples: number of samples to use to generate plot
"""
sampler = MarkovJumpHMC(distribution=distr, **kwargs)
samples = sampler.sample(nsamples)
with sns.axes_style("white"):
g = sns.jointplot(samples[0], samples[1], kind='kde', stat_func=None)
return g
def gauss_2d(nsamples=1000):
"""
Another simple test plot
1d gaussian sampled from each sampler visualized as a joint 2d gaussian
"""
gaussian = TestGaussian(ndims=1)
control = HMCBase(distribution=gaussian)
experimental = MarkovJumpHMC(distribution=gaussian, resample=False)
with sns.axes_style("white"):
sns.jointplot(control.sample(nsamples)[0],
experimental.sample(nsamples)[0],
kind='hex',
stat_func=None)
def hist_1d(distr, nsamples=1000, nbins=250, control=True, resample=True):
"""
plots a 1d histogram from each sampler
distr is (an unitialized) class from distributions
"""
distribution = distr(ndims=1)
control_smp = HMCBase(distribution=distribution, epsilon=1)
experimental_smp = MarkovJumpHMC(distribution=distribution, resample=resample, epsilon=1)
if control:
plt.hist(control_smp.sample(nsamples)[0], nbins, normed=True, label="Standard HMCBase", alpha=.5)
plt.hist(experimental_smp.sample(nsamples)[0], nbins, normed=True, label="Continuous-time HMCBase",alpha=.5)
plt.legend()
def hist_2d(distr, nsamples, **kwargs):
"""
Plots a 2d hexbinned histogram of distribution
Args:
distr: Distribution object
nsamples: number of samples to use to generate plot
"""
sampler = MarkovJumpHMC(distribution=distr, **kwargs)
samples = sampler.sample(nsamples)
with sns.axes_style("white"):
g = sns.jointplot(samples[0], samples[1], kind='kde', stat_func=None)
return g
def generate_figure_samples(samples_per_frame, n_frames, burnin = int(1e4)):
""" Generates the figure
:param samples_per_frame: number of sample steps between each frame
:param n_frames: number of frames to draw
:returns: None
:rtype: None
"""
n_samples = samples_per_frame * n_frames
ndims = 36
nbasis = 72
rand_val = rand(ndims,nbasis/2,density=0.25)
W = np.concatenate([rand_val.toarray(), -rand_val.toarray()],axis=1)
logalpha = np.random.randn(nbasis, 1)
poe = ProductOfT(nbatch=1, W=W, logalpha=logalpha)
## NUTS uses a different number of grad evals for each update step!!
## makes it very hard to compare against others w/ same number of update steps
# # NUTS
# print "NUTS"
# nuts_init = poe.Xinit[:, 0]
# nuts_samples = nuts6(poe.reset(), n_samples, nuts_burnin, nuts_init)[0]
# nuts_frames = [nuts_samples[f_idx * samples_per_frame, :] for f_idx in xrange(0, n_frames)]
# x_init = nuts_samples[0, :].reshape(ndims, 1)
## burnin
print "MJHMC burnin"
x_init = poe.Xinit #[:, [0]]
mjhmc = MarkovJumpHMC(distribution=poe.reset(), **mjhmc_params)
mjhmc.state = HMCState(x_init.copy(), mjhmc)
mjhmc_samples = mjhmc.sample(burnin)
print mjhmc_samples.shape
x_init = mjhmc_samples[:, [0]]
# control HMC
print "Control"
hmc = ControlHMC(distribution=poe.reset(), **control_params)
hmc.state = HMCState(x_init.copy(), hmc)
hmc_samples = hmc.sample(n_samples)
hmc_frames = [hmc_samples[:, f_idx * samples_per_frame].copy() for f_idx in xrange(0, n_frames)]
# MJHMC
print "MJHMC"
mjhmc = MarkovJumpHMC(distribution=poe.reset(), resample=False, **mjhmc_params)
mjhmc.state = HMCState(x_init.copy(), mjhmc)
mjhmc_samples = mjhmc.sample(n_samples)
mjhmc_frames = [mjhmc_samples[:, f_idx * samples_per_frame].copy() for f_idx in xrange(0, n_frames)]
print mjhmc.r_count, hmc.r_count
print mjhmc.l_count, hmc.l_count
print mjhmc.f_count, hmc.f_count
print mjhmc.fl_count, hmc.fl_count
frames = [mjhmc_frames, hmc_frames]
names = ['MJHMC', 'ControlHMC']
frame_grads = [f_idx * samples_per_frame for f_idx in xrange(0, n_frames)]
return frames, names, frame_grads
def gauss_2d(nsamples=1000):
"""
Another simple test plot
1d gaussian sampled from each sampler visualized as a joint 2d gaussian
"""
gaussian = TestGaussian(ndims=1)
control = HMCBase(distribution=gaussian)
experimental = MarkovJumpHMC(distribution=gaussian, resample=False)
with sns.axes_style("white"):
sns.jointplot(
control.sample(nsamples)[0],
experimental.sample(nsamples)[0],
kind='hex',
stat_func=None)
def generate_initialization(distribution):
""" Run mjhmc for BURN_IN_STEPS on distribution, generating a fair set of initial states
:param distribution: Distribution object. Must have nbatch == MAX_N_PARTICLES
:returns: a set of fair initial states and an estimate of the variance for emc and true both
:rtype: tuple: (array of shape (distribution.ndims, MAX_N_PARTICLES), float, float)
"""
print(
'Generating fair initialization for {} by burning in {} steps'.format(
type(distribution).__name__, BURN_IN_STEPS))
assert BURN_IN_STEPS > VAR_STEPS
# must rebuild graph to nbatch=MAX_N_PARTICLES
if distribution.backend == 'tensorflow':
distribution.build_graph()
mjhmc = MarkovJumpHMC(distribution=distribution, resample=False)
for _ in xrange(BURN_IN_STEPS - VAR_STEPS):
mjhmc.sampling_iteration()
assert mjhmc.resample == False
emc_var_estimate, mjhmc = online_variance(mjhmc, distribution)
# we discard v since p(x,v) = p(x)p(v)
mjhmc_endpt = mjhmc.state.copy().X
# otherwise will go into recursive loop
distribution.mjhmc = False
try:
distribution.gen_init_X()
except NotImplementedError:
print("No explicit init method found, using mjhmc endpoint")
distribution.E_count = 0
distribution.dEdX_count = 0
control = ControlHMC(distribution=distribution)
for _ in xrange(BURN_IN_STEPS - VAR_STEPS):
control.sampling_iteration()
true_var_estimate, control = online_variance(control, distribution)
control_endpt = control.state.copy().X
return mjhmc_endpt, emc_var_estimate, true_var_estimate, control_endpt
def generate_initialization(distribution):
""" Run mjhmc for BURN_IN_STEPS on distribution, generating a fair set of initial states
:param distribution: Distribution object. Must have nbatch == MAX_N_PARTICLES
:returns: a set of fair initial states and an estimate of the variance for emc and true both
:rtype: tuple: (array of shape (distribution.ndims, MAX_N_PARTICLES), float, float)
"""
print('Generating fair initialization for {} by burning in {} steps'.format(
type(distribution).__name__, BURN_IN_STEPS))
assert BURN_IN_STEPS > VAR_STEPS
# must rebuild graph to nbatch=MAX_N_PARTICLES
if distribution.backend == 'tensorflow':
distribution.build_graph()
mjhmc = MarkovJumpHMC(distribution=distribution, resample=False)
for _ in xrange(BURN_IN_STEPS - VAR_STEPS):
mjhmc.sampling_iteration()
assert mjhmc.resample == False
emc_var_estimate, mjhmc = online_variance(mjhmc, distribution)
# we discard v since p(x,v) = p(x)p(v)
mjhmc_endpt = mjhmc.state.copy().X
# otherwise will go into recursive loop
distribution.mjhmc = False
try:
distribution.gen_init_X()
except NotImplementedError:
print("No explicit init method found, using mjhmc endpoint")
distribution.E_count = 0
distribution.dEdX_count = 0
control = ControlHMC(distribution=distribution)
for _ in xrange(BURN_IN_STEPS - VAR_STEPS):
control.sampling_iteration()
true_var_estimate, control = online_variance(control, distribution)
control_endpt = control.state.copy().X
return mjhmc_endpt, emc_var_estimate, true_var_estimate, control_endpt
def jump_plot(distribution, nsamples=100, **kwargs):
"""
Plots samples drawn from distribution with dwelling time on the x-axis
and the sample value on the y-axis
1D only
"""
distr = distribution(ndims=1, nbatch=1, **kwargs)
# sampler = MarkovJumpHMC(np.array([0]).reshape(1,1),
sampler = MarkovJumpHMC(distr.Xinit,
distr.E,
distr.dEdX,
epsilon=.3,
beta=.2,
num_leapfrog_steps=5)
x_t = []
d_t = []
transitions = []
last_L_count, last_F_count, last_R_count = 0, 0, 0
for idx in xrange(nsamples):
sampler.sampling_iteration()
x_t.append(sampler.state.X[0, 0])
d_t.append(sampler.dwelling_times[0])
if sampler.L_count - last_L_count == 1:
transitions.append("L")
last_L_count += 1
elif sampler.F_count - last_F_count == 1:
transitions.append("F")
last_F_count += 1
elif sampler.R_count - last_R_count == 1:
transitions.append("R")
last_R_count += 1
t = np.cumsum(d_t)
plt.scatter(t, x_t)
t = np.array(t).reshape(len(t), 1)
x_t = np.array(x_t).reshape(len(x_t), 1)
transitions = np.array(transitions).reshape(len(transitions), 1)
data = np.concatenate((x_t, t, transitions), axis=1)
return pd.DataFrame(data, columns=['x', 't', 'transitions'])
def generate_figure_samples(samples_per_frame, n_frames, burnin=int(1e4)):
""" Generates the figure
:param samples_per_frame: number of sample steps between each frame
:param n_frames: number of frames to draw
:returns: None
:rtype: None
"""
n_samples = samples_per_frame * n_frames
ndims = 36
nbasis = 72
rand_val = rand(ndims, nbasis / 2, density=0.25)
W = np.concatenate([rand_val.toarray(), -rand_val.toarray()], axis=1)
logalpha = np.random.randn(nbasis, 1)
poe = ProductOfT(nbatch=1, W=W, logalpha=logalpha)
## NUTS uses a different number of grad evals for each update step!!
## makes it very hard to compare against others w/ same number of update steps
# # NUTS
# print "NUTS"
# nuts_init = poe.Xinit[:, 0]
# nuts_samples = nuts6(poe.reset(), n_samples, nuts_burnin, nuts_init)[0]
# nuts_frames = [nuts_samples[f_idx * samples_per_frame, :] for f_idx in xrange(0, n_frames)]
# x_init = nuts_samples[0, :].reshape(ndims, 1)
## burnin
print "MJHMC burnin"
x_init = poe.Xinit # [:, [0]]
mjhmc = MarkovJumpHMC(distribution=poe.reset(), **mjhmc_params)
mjhmc.state = HMCState(x_init.copy(), mjhmc)
mjhmc_samples = mjhmc.sample(burnin)
print mjhmc_samples.shape
x_init = mjhmc_samples[:, [0]]
# control HMC
print "Control"
hmc = ControlHMC(distribution=poe.reset(), **control_params)
hmc.state = HMCState(x_init.copy(), hmc)
hmc_samples = hmc.sample(n_samples)
hmc_frames = [hmc_samples[:, f_idx * samples_per_frame].copy() for f_idx in xrange(0, n_frames)]
# MJHMC
print "MJHMC"
mjhmc = MarkovJumpHMC(distribution=poe.reset(), resample=False, **mjhmc_params)
mjhmc.state = HMCState(x_init.copy(), mjhmc)
mjhmc_samples = mjhmc.sample(n_samples)
mjhmc_frames = [mjhmc_samples[:, f_idx * samples_per_frame].copy() for f_idx in xrange(0, n_frames)]
print mjhmc.r_count, hmc.r_count
print mjhmc.l_count, hmc.l_count
print mjhmc.f_count, hmc.f_count
print mjhmc.fl_count, hmc.fl_count
frames = [mjhmc_frames, hmc_frames]
names = ["MJHMC", "ControlHMC"]
frame_grads = [f_idx * samples_per_frame for f_idx in xrange(0, n_frames)]
return frames, names, frame_grads
def hist_1d(distr, nsamples=1000, nbins=250, control=True, resample=True):
"""
plots a 1d histogram from each sampler
distr is (an unitialized) class from distributions
"""
distribution = distr(ndims=1)
control_smp = HMCBase(distribution=distribution, epsilon=1)
experimental_smp = MarkovJumpHMC(distribution=distribution,
resample=resample,
epsilon=1)
if control:
plt.hist(control_smp.sample(nsamples)[0],
nbins,
normed=True,
label="Standard HMCBase",
alpha=.5)
plt.hist(experimental_smp.sample(nsamples)[0],
nbins,
normed=True,
label="Continuous-time HMCBase",
alpha=.5)
plt.legend()
