def compute_trajectory(self, random_start_state=None):
logger.debug("Entering")
if random_start_state is not None:
np.random.set_state(random_start_state)
else:
random_start_state = np.random.get_state()
# momentum
L_p = np.linalg.cholesky(np.eye(self.D) * self.sigma_p)
self.logp = lambda x: log_gaussian_pdf(
x, Sigma=L_p, compute_grad=False, is_cholesky=True)
self.dlogp = lambda x: log_gaussian_pdf(
x, Sigma=L_p, compute_grad=True, is_cholesky=True)
self.p_sample = lambda: sample_gaussian(
N=1, mu=np.zeros(self.D), Sigma=L_p, is_cholesky=True)[0]
# set up target and momentum densities and gradients
self.set_up()
dlogq_est = self.update_density_estimate()
# random number of steps?
if self.max_steps is not None:
steps = np.random.randint(self.num_steps, self.max_steps + 1)
else:
steps = self.num_steps
logger.info("Simulating trajectory for at least L=%d steps of size %.2f" % \
(self.num_steps, self.step_size))
# starting state
p0 = self.p_sample()
q0 = self.q_sample()
Qs, Ps = leapfrog(q0, self.dlogq, p0, self.dlogp, self.step_size,
steps)
# run second integrator for same amount of steps
steps_taken = len(Qs)
Qs_est, Ps_est = leapfrog(q0, dlogq_est, p0, self.dlogp,
self.step_size, steps_taken)
logger.info("%d steps taken" % steps_taken)
logger.info("Computing average acceptance probabilities")
log_acc = compute_log_accept_pr(q0, p0, Qs, Ps, self.logq, self.logp)
log_acc_est = compute_log_accept_pr(q0, p0, Qs_est, Ps_est, self.logq,
self.logp)
acc_mean = np.mean(np.exp(log_acc))
acc_est_mean = np.mean(np.exp(log_acc_est))
idx09 = int(len(log_acc) * 0.9)
acc_mean10 = np.mean(np.exp(log_acc[idx09:]))
acc_est_mean10 = np.mean(np.exp(log_acc_est[idx09:]))
logger.info("Computing average volumes")
log_det = compute_log_det_trajectory(Qs, Ps)
log_det_est = compute_log_det_trajectory(Qs_est, Ps_est)
logger.info("Average acceptance prob: %.2f, %.2f" %
(acc_mean, acc_est_mean))
logger.info("Average acceptance prob (last 10 percent): %.2f, %.2f" %
(acc_mean10, acc_est_mean10))
logger.info("Log-determinant: %.2f, %.2f" % (log_det, log_det_est))
logger.debug("Leaving")
return acc_mean, acc_est_mean, log_det, log_det_est, steps_taken, random_start_state
def compute_trajectory(self, random_start_state=None):
logger.debug("Entering")
if random_start_state is not None:
np.random.set_state(random_start_state)
else:
random_start_state = np.random.get_state()
# momentum
L_p = np.linalg.cholesky(np.eye(self.D) * self.sigma_p)
self.logp = lambda x: log_gaussian_pdf(
x, Sigma=L_p, compute_grad=False, is_cholesky=True)
self.dlogp = lambda x: log_gaussian_pdf(
x, Sigma=L_p, compute_grad=True, is_cholesky=True)
self.p_sample = lambda: sample_gaussian(
N=1, mu=np.zeros(self.D), Sigma=L_p, is_cholesky=True)[0]
self.p_sample = lambda: sample_gaussian(
N=1, mu=np.zeros(self.D), Sigma=L_p, is_cholesky=True)[0]
# set up target and momentum densities and gradients
self.set_up()
logger.info("Learning kernel bandwidth")
sigma = select_sigma_grid(self.Z, lmbda=self.lmbda, log2_sigma_max=15)
logger.info("Using lmbda=%.2f, sigma: %.2f" % (self.lmbda, sigma))
logger.info("Computing kernel matrix")
K = gaussian_kernel(self.Z, sigma=sigma)
logger.info("Estimate density in RKHS")
b = _compute_b_sym(self.Z, K, sigma)
C = _compute_C_sym(self.Z, K, sigma)
a = score_matching_sym(self.Z, sigma, self.lmbda, K, b, C)
# logger.info("Computing objective function")
# J = _objective_sym(Z, sigma, self.lmbda, a, K, b, C)
# J_xval = np.mean(xvalidate(Z, 5, sigma, self.lmbda, K))
# logger.info("N=%d, sigma: %.2f, lambda: %.2f, J(a)=%.2f, XJ(a)=%.2f" % \
# (self.N, sigma, self.lmbda, J, J_xval))
kernel_grad = lambda x, X=None: gaussian_kernel_grad(x, X, sigma)
dlogq_est = lambda x: log_pdf_estimate_grad(x, a, self.Z, kernel_grad)
logger.info("Simulating trajectory for L=%d steps of size %.2f" % \
(self.num_steps, self.step_size))
# starting state
p0 = self.p_sample()
q0 = self.q_sample()
Qs, Ps = leapfrog(q0, self.dlogq, p0, self.dlogp, self.step_size,
self.num_steps, self.max_steps)
# run second integrator for same amount of steps
steps_taken = len(Qs)
logger.info("%d steps taken" % steps_taken)
Qs_est, Ps_est = leapfrog(q0, dlogq_est, p0, self.dlogp,
self.step_size, steps_taken)
logger.info("Computing average acceptance probabilities")
log_acc = compute_log_accept_pr(q0, p0, Qs, Ps, self.logq, self.logp)
log_acc_est = compute_log_accept_pr(q0, p0, Qs_est, Ps_est, self.logq,
self.logp)
acc_mean = np.mean(np.exp(log_acc))
acc_est_mean = np.mean(np.exp(log_acc_est))
logger.info("Computing average volumes")
log_det = compute_log_det_trajectory(Qs, Ps)
log_det_est = compute_log_det_trajectory(Qs_est, Ps_est)
logger.info("Average acceptance prob: %.2f, %.2f" %
(acc_mean, acc_est_mean))
logger.info("Log-determinant: %.2f, %.2f" % (log_det, log_det_est))
logger.debug("Leaving")
return acc_mean, acc_est_mean, log_det, log_det_est, steps_taken, random_start_state
# initial
samples[0] = q_current
samples_est[0] = q_current_est
# run MCMC
for i in np.arange(1, num_iterations):
print("%d/%d" % (i + 1, num_iterations))
# sample momentum
p = p_sample()
momentums[i] = p
# simulate Hamiltonian flow, use last point as proposal
num_steps = np.random.randint(num_steps_min, num_steps_max + 1)
step_size = np.random.rand() * (step_size_max - step_size_min) + step_size_min
Qs, Ps = leapfrog(q_current, dlogq, p, dlogp, step_size, num_steps)
Qs_est, Ps_est = leapfrog(q_current_est, dlogq_est, p, dlogp, step_size, num_steps)
proposals[i] = Qs[-1]
proposals_est[i] = Qs_est[-1]
# compute acceptance probability
acc_prob[i] = np.exp(compute_log_accept_pr_single(q_current, p, Qs[-1], Ps[-1], logq, logp))
acc_prob_est[i] = np.exp(compute_log_accept_pr_single(q_current_est, p, Qs_est[-1], Ps_est[-1], logq, logp))
if False:
# visualise trajectories and acceptance probability along
plot_array(Xs, Ys, np.exp(G), plot_contour=False)
plot_2d_trajectory(Qs, "r-")
plot_2d_trajectory(Qs_est, "b-")
plt.show()
def propose(self, current, current_log_pdf, samples, accepted):
# random variables from a fixed random stream without modifying the current one
rnd_state = np.random.get_state()
np.random.set_state(self.hmc_rnd_state)
if current_log_pdf is None:
current_log_pdf = self.orig_target.log_pdf(current)
# sample momentum and leapfrog parameters
p0 = self.momentum.sample()
num_steps = np.random.randint(self.num_steps_min, self.num_steps_max + 1)
step_size = np.random.rand() * (self.step_size_max - self.step_size_min) + self.step_size_min
# restore random state
self.hmc_rnd_state = np.random.get_state()
np.random.set_state(rnd_state)
logger.debug("Simulating Hamiltonian flow")
Qs, Ps = leapfrog(current, self.target.grad, p0, self.momentum.grad, step_size, num_steps)
q=Qs[-1]
p=Ps[-1]
logger.debug("Momentum start: %s" % str(p0))
logger.debug("Momentum end: %s" % str(p))
# compute acceptance probability, extracting log_pdf of q
p0_log_pdf = self.momentum.log_pdf(p0)
p_log_pdf = self.momentum.log_pdf(p)
# use a function call to be able to overload it for KMC
acc_prob, log_pdf_q = self.accept_prob_log_pdf(current, q, p0_log_pdf, p_log_pdf, current_log_pdf, samples)
if True and (len(samples) % 100) ==0:
logger.debug("Plotting")
import matplotlib.pyplot as plt
res = 50
Xs_q = np.linspace(-4,4, res)
Ys_q = np.linspace(-4,4, res)
# evaluate density and estimate
D1=0
D2=1
def dummy_grad(X_2d):
theta = current.copy()
# theta = np.mean(self.Z, 0)
theta[D1]=X_2d[0]
theta[D2]=X_2d[1]
return self.target.grad(theta)
def dummy(X_2d):
theta = current.copy()
# theta = np.mean(self.Z, 0)
theta[D1]=X_2d[0]
theta[D2]=X_2d[1]
return self.target.log_pdf(theta)
# plt.figure()
# G = evaluate_density_grid(Xs_q, Ys_q, dummy)
# plot_array(Xs_q, Ys_q, G)
# plt.plot(self.Z[:,D1], self.Z[:,D2], '.')
# plt.plot(Qs[:,D1], Qs[:,D2], 'r-')
# plt.plot(samples[:,D1], samples[:,D2], 'm-')
# plt.plot(current[D1], current[D2], 'b*', markersize=15)
# plt.plot(Qs[-1,D1], Qs[-1,D2], 'r*', markersize=15)
plt.figure()
G_norm, U_q, V, X, Y = evaluate_gradient_grid(Xs_q, Ys_q, dummy_grad)
plot_array(Xs_q, Ys_q, G_norm)
plt.plot(self.Z[:,D1], self.Z[:,D2], '.')
plt.plot(Qs[:,D1], Qs[:,D2], 'r-')
plt.plot(samples[:,D1], samples[:,D2], 'm-')
plt.plot(current[D1], current[D2], 'b*', markersize=15)
plt.plot(Qs[-1,D1], Qs[-1,D2], 'r*', markersize=15)
plt.quiver(X, Y, U_q, V, color='m')
# plt.figure()
# plt.plot(Ps[:,D1], Ps[:,D2], 'r-')
# plt.plot(p0[D1], p0[D2], 'b*', markersize=15)
# plt.plot(Ps[-1,D1], Ps[-1,D2], 'r*', markersize=15)
# plt.title('momentum')
acc_probs = np.exp(compute_log_accept_pr(current, p0, Qs, Ps, self.orig_target.log_pdf, self.momentum.log_pdf))
H_ratios = np.exp(compute_log_accept_pr(current, p0, Qs, Ps, self.target.log_pdf, self.momentum.log_pdf))
target_ratio = [np.min([1,np.exp(self.orig_target.log_pdf(x)-current_log_pdf)]) for x in Qs]
momentum_ratio = [np.min([1,np.exp(self.momentum.log_pdf(x)-p0_log_pdf)]) for x in Ps]
target_log_pdf = np.exp(np.array([self.orig_target.log_pdf(x) for x in Qs]))
# #
# plt.figure(figsize=(12,4))
# plt.subplot(151)
# plt.plot(acc_probs)
# plt.plot([0, len(acc_probs)], [acc_probs.mean(), acc_probs.mean()])
# plt.title("acc_probs")
# plt.subplot(152)
# plt.plot(target_ratio)
# plt.title("target_ratio")
# plt.subplot(153)
# plt.plot(momentum_ratio)
# plt.title("momentum_ratio")
# plt.subplot(154)
# plt.plot(H_ratios)
# plt.title("H_ratios")
# plt.subplot(155)
# plt.plot(target_log_pdf)
# plt.title("target_log_pdf")
plt.show()
return q, acc_prob, log_pdf_q
def plot_kamiltonian_dnyamics(q0, p0, logq, dlogq, logq_est, dlogq_est,
logp, dlogp, Z=None, num_steps=500, step_size=.1,
Xs_q=None, Ys_q=None, Xs_p=None, Ys_p=None,
plot_dlogq=False, plot_H_or_acc=True):
D = len(q0)
# compute and plot log-density, if D==2
plt.figure(figsize=(12, 12))
if D is 2:
if Xs_q is None:
Xs_q = np.linspace(-3, 3)
if Ys_q is None:
Ys_q = np.linspace(-3, 3)
if Xs_p is None:
Xs_p = np.linspace(-3, 3)
if Ys_p is None:
Ys_p = np.linspace(-3, 3)
if plot_dlogq:
G = evaluate_density_grad_grid(Xs_q, Ys_q, dlogq)
G_est = evaluate_density_grad_grid(Xs_q, Ys_q, dlogq_est)
else:
G = evaluate_density_grid(Xs_q, Ys_q, logq)
G_est = evaluate_density_grid(Xs_q, Ys_q, logq_est)
M = evaluate_density_grid(Xs_p, Ys_p, logp)
M_est = evaluate_density_grid(Xs_p, Ys_p, logp)
plt.subplot(321)
if not plot_dlogq:
plot_array(Xs_q, Ys_q, np.exp(G))
else:
plot_array(Xs_q, Ys_q, G)
if Z is not None:
plt.plot(Z[:, 0], Z[:, 1], 'bx')
plt.subplot(322)
plot_array(Xs_q, Ys_q, np.exp(G_est))
if Z is not None:
plt.plot(Z[:, 0], Z[:, 1], 'bx')
plt.subplot(323)
plot_array(Xs_p, Ys_p, np.exp(M))
plt.subplot(324)
plot_array(Xs_p, Ys_p, np.exp(M_est))
Qs, Ps = leapfrog(q0, dlogq, p0, dlogp, step_size, num_steps)
Qs_est, Ps_est = leapfrog(q0, dlogq_est, p0, dlogp, step_size, num_steps)
Hs = compute_hamiltonian(Qs, Ps, logq, logp)
Hs_est = compute_hamiltonian(Qs_est, Ps_est, logq, logp)
log_acc = compute_log_accept_pr(q0, p0, Qs, Ps, logq, logp)
log_acc_est = compute_log_accept_pr(q0, p0, Qs_est, Ps_est, logq, logp)
acc_mean = np.mean(np.exp(log_acc))
acc_est_mean = np.mean(np.exp(log_acc_est))
logger.info("HMC acceptance prob: %.2f" % acc_mean)
logger.info("KMC acceptance prob: %.2f" % acc_est_mean)
spread = compute_log_det_trajectory(Qs, Ps)
spread_est = compute_log_det_trajectory(Qs_est, Ps_est)
logger.info("HMC spread: %.2f" % (spread))
logger.info("KMC spread: %.2f" % (spread_est))
plt.subplot(321)
plot_2d_trajectory(Qs)
plt.title("True density")
plt.subplot(322)
plot_2d_trajectory(Qs_est)
plt.title("Estimated density")
plt.subplot(323)
plt.title("Momentum")
plot_2d_trajectory(Ps)
plt.subplot(324)
plt.title("Momentum")
plot_2d_trajectory(Ps_est)
if plot_H_or_acc:
ylim = [np.min([Hs.min(), Hs_est.min()]),
np.max([Hs.max(), Hs_est.max()])]
plt.subplot(325)
plt.title("Hamiltonian")
plt.plot(Hs)
plt.ylim(ylim)
plt.gca().get_yaxis().get_major_formatter().set_useOffset(False)
plt.subplot(326)
plt.title("Hamiltonian")
plt.plot(Hs_est)
plt.ylim(ylim)
plt.gca().get_yaxis().get_major_formatter().set_useOffset(False)
else:
plt.subplot(325)
plt.title("Acceptance prob.")
plt.plot(np.exp(log_acc))
plt.plot([0,len(log_acc)], [acc_mean, acc_mean], "r")
plt.gca().get_yaxis().get_major_formatter().set_useOffset(False)
plt.subplot(326)
plt.title("Acceptance prob.")
plt.plot(np.exp(log_acc_est))
plt.plot([0,len(log_acc_est)], [acc_est_mean, acc_est_mean], "r")
plt.gca().get_yaxis().get_major_formatter().set_useOffset(False)
plt.tight_layout()
# plotting grid
res = 200
Xs_q = np.linspace(-3, 3, res)
Ys_q = np.linspace(-3, 3, res)
Xs_p = np.linspace(-1, 1, res)
Ys_p = np.linspace(-1, 1, res)
# evaluate density and estimate
G = evaluate_density_grid(Xs_q, Ys_q, logq)
G_est = evaluate_density_grid(Xs_q, Ys_q, logq_est)
# evaluate momentum, which is the same for both
M = evaluate_density_grid(Xs_p, Ys_p, logp)
# simulate true and approximate Hamiltonian
Qs, Ps = leapfrog(q0, dlogq, p0, dlogp, step_size, num_steps)
Qs_est, Ps_est = leapfrog(q0, dlogq_est, p0, dlogp, step_size, num_steps)
Hs = compute_hamiltonian(Qs, Ps, logq, logp)
Hs_est = compute_hamiltonian(Qs_est, Ps_est, logq, logp)
# compute acceptance probabilities
log_acc = compute_log_accept_pr(q0, p0, Qs, Ps, logq, logp)
log_acc_est = compute_log_accept_pr(q0, p0, Qs_est, Ps_est, logq, logp)
# normalise Hamiltonians
Hs -= Hs.mean()
Hs_est -= Hs_est.mean()
plt.figure()
plot_array(Xs_q, Ys_q, np.exp(G))
plot_2d_trajectory(Qs)