def set_up(self): # place a gamma on the Eigenvalues of a Gaussian covariance EVs = np.random.gamma(shape=1, size=self.D) # random orthogonal matrix to rotate Q = qmult(np.eye(self.D)) Sigma = Q.T.dot(np.diag(EVs)).dot(Q) # Cholesky of random covariance L = np.linalg.cholesky(Sigma) # target density self.dlogq = lambda x: log_gaussian_pdf( x, Sigma=L, is_cholesky=True, compute_grad=True) self.logq = lambda x: log_gaussian_pdf( x, Sigma=L, is_cholesky=True, compute_grad=False) # starting state self.q_sample = lambda: sample_gaussian( N=1, mu=np.zeros(self.D), Sigma=L, is_cholesky=True)[0] logger.info("N=%d, D=%d" % (self.N, self.D)) self.Z = sample_gaussian(self.N, mu=np.zeros(self.D), Sigma=L, is_cholesky=True)
def set_up(self): L = np.linalg.cholesky(np.eye(self.D) * self.sigma_q) # target density self.dlogq = lambda x: log_gaussian_pdf( x, Sigma=L, is_cholesky=True, compute_grad=True) self.logq = lambda x: log_gaussian_pdf( x, Sigma=L, is_cholesky=True, compute_grad=False) # starting state self.q_sample = lambda: sample_gaussian( N=1, mu=np.zeros(self.D), Sigma=L, is_cholesky=True)[0] logger.info("N=%d, D=%d" % (self.N, self.D)) self.Z = sample_gaussian(self.N, mu=np.zeros(self.D), Sigma=L, is_cholesky=True)
def leapfrog_friction_habc_no_storing(c, V, q, dlogq, p, dlogp, step_size=0.3, num_steps=1): logger.debug("Entering") """ MATLAB code by Chen: function [ newx ] = sghmc( U, gradU, m, dt, nstep, x, C, V ) %% SGHMC using gradU, for nstep, starting at position x p = randn( size(x) ) * sqrt( m ); B = 0.5 * V * dt; D = sqrt( 2 * (C-B) * dt ); for i = 1 : nstep p = p - gradU( x ) * dt - p * C * dt + randn(1)*D; x = x + p./m * dt; end newx = x; end """ # friction term (as in HABC) D = len(q) B = 0.5 * V * step_size C = np.eye(D) * c + V L_friction = np.linalg.cholesky(2 * step_size * (C - B)) zeros_D = np.zeros(D) # create copy of state p = np.array(p.copy()) q = np.array(q.copy()) # alternate full momentum and variable updates for _ in range(num_steps): friction = sample_gaussian(N=1, mu=zeros_D, Sigma=L_friction, is_cholesky=True)[0] # just like normal momentum update but with friction p = p - step_size * -dlogq(q) - step_size * C.dot(-dlogp(p)) + friction # normal position update q = q + step_size * -dlogp(p) logger.debug("Leaving") return q, p
print("N=%d, sigma: %.2f, lambda: %.2f, J(a)=%.2f, XJ(a)=%.2f" % \ (N, sigma, lmbda, J, J_xval)) kernel = lambda X, Y=None: gaussian_kernel(X, Y, sigma=sigma) kernel_grad = lambda x, X=None: gaussian_kernel_grad(x, X, sigma) logq_est = lambda x: log_pdf_estimate(x, a, Z, kernel) dlogq_est = lambda x: log_pdf_estimate_grad(x, a, Z, kernel_grad) # momentum Sigma_p = np.eye(D) * .1 L_p = np.linalg.cholesky(Sigma_p) logp = lambda x: log_gaussian_pdf( x, Sigma=L_p, compute_grad=False, is_cholesky=True) dlogp = lambda x: log_gaussian_pdf( x, Sigma=L_p, compute_grad=True, is_cholesky=True) p_sample = lambda: sample_gaussian( N=1, mu=np.zeros(D), Sigma=L_p, is_cholesky=True)[0] # starting state p0 = p_sample() q0 = np.zeros(D) q0[:2] = np.array([0, -3]) # parameters num_steps = 1500 step_size = .1 Xs_q = np.linspace(-20, 20) Ys_q = np.linspace(-10, 10) Xs_p = np.linspace(-1, 1) Ys_p = np.linspace(-1, 1)
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
if __name__ == "__main__": D = 2 # true target log density Sigma = np.diag(np.linspace(0.01, 1, D)) Sigma[:2, :2] = np.array([[1, .95], [.95, 1]]) Sigma = np.eye(D) L = np.linalg.cholesky(Sigma) # estimate density in rkhs N = 200 mu = np.zeros(D) np.random.seed(0) Z = sample_gaussian(N, mu, Sigma=L, is_cholesky=True) # print np.sum(Z) * np.std(Z) * np.sum(Z**2) * np.std(Z**2) resolution = 10 log2_sigmas = np.linspace(-3, 3, resolution) log2_lambdas = np.linspace(-15, 3, resolution) Js = np.zeros((len(log2_sigmas), len(log2_lambdas))) Js_std = np.zeros(Js.shape) pool = Pool() chunksize = 1 for ind, res in enumerate( pool.imap(fun, product(log2_sigmas, log2_lambdas)), chunksize): Js.flat[ind - 1], Js_std.flat[ind - 1] = res im = plt.pcolor(log2_lambdas, log2_sigmas, np.log2(Js - Js.min() + 1))