diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index 0c42de6bf9..9e22a1639a 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/pymc3/step_methods/hmc/base_hmc.py @@ -1,5 +1,5 @@ from ..arraystep import ArrayStepShared -from .trajectory import get_theano_hamiltonian_functions +from .trajectory import get_theano_hamiltonian_functions, get_theano_hamiltonian_manifold_functions from pymc3.tuning import guess_scaling from pymc3.model import modelcontext, Point @@ -13,7 +13,7 @@ class BaseHMC(ArrayStepShared): def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, model=None, blocked=True, use_single_leapfrog=False, - potential=None, integrator="leapfrog", **theano_kwargs): + potential=None, integrator="leapfrog", manifold=False, **theano_kwargs): """Superclass to implement Hamiltonian/hybrid monte carlo Parameters @@ -60,7 +60,11 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, if theano_kwargs is None: theano_kwargs = {} - self.H, self.compute_energy, self.compute_velocity, self.leapfrog, self.dlogp = get_theano_hamiltonian_functions( - vars, shared, model.logpt, self.potential, use_single_leapfrog, integrator, **theano_kwargs) + if manifold: + self.H, self.compute_energy, self.compute_velocity, self.leapfrog, self.dlogp = get_theano_hamiltonian_manifold_functions( + vars, shared, model.logpt, self.potential, **theano_kwargs) + else: + self.H, self.compute_energy, self.compute_velocity, self.leapfrog, self.dlogp = get_theano_hamiltonian_functions( + vars, shared, model.logpt, self.potential, use_single_leapfrog, integrator, **theano_kwargs) super(BaseHMC, self).__init__(vars, shared, blocked=blocked) diff --git a/pymc3/step_methods/hmc/rmhmc.py b/pymc3/step_methods/hmc/rmhmc.py new file mode 100644 index 0000000000..13aeb9b4a8 --- /dev/null +++ b/pymc3/step_methods/hmc/rmhmc.py @@ -0,0 +1,63 @@ +from ..arraystep import metrop_select, Competence +from .base_hmc import BaseHMC +from pymc3.vartypes import discrete_types +from pymc3.theanof import floatX +import numpy as np + + +__all__ = ['RiemannianManifoldHMC'] + +def unif(step_size, elow=.85, ehigh=1.15): + return np.random.uniform(elow, ehigh) * step_size + +class RiemannianManifoldHMC(BaseHMC): + name = 'rmhmc' + default_blocked = True + + def __init__(self, vars=None, path_length=2., step_rand=unif, **kwargs): + """ + Parameters + ---------- + vars : list of theano variables + path_length : float, default=2 + total length to travel + step_rand : function float -> float, default=unif + A function which takes the step size and returns an new one used to + randomize the step size at each iteration. + step_scale : float, default=0.25 + Initial size of steps to take, automatically scaled down + by 1/n**(1/4). + scaling : array_like, ndim = {1,2} + The inverse mass, or precision matrix. One dimensional arrays are + interpreted as diagonal matrices. If `is_cov` is set to True, + this will be interpreded as the mass or covariance matrix. + is_cov : bool, default=False + Treat the scaling as mass or covariance matrix. + potential : Potential, optional + An object that represents the Hamiltonian with methods `velocity`, + `energy`, and `random` methods. It can be specified instead + of the scaling matrix. + model : pymc3.Model + The model + **kwargs : passed to BaseHMC + """ + super(RiemannianManifoldHMC, self).__init__(vars, manifold=True, **kwargs) + self.path_length = path_length + self.step_rand = step_rand + + def astep(self, q0): + e = floatX(self.step_rand(self.step_size)) + n_steps = np.array(self.path_length / e, dtype='int32') + q = q0 + p = self.H.pot.random() # initialize momentum + initial_energy = self.compute_energy(q, p) + q, p, current_energy = self.leapfrog(q, p, e, n_steps) + energy_change = initial_energy - current_energy + return metrop_select(energy_change, q, q0)[0] + + @staticmethod + def competence(var): + if var.dtype in discrete_types: + return Competence.INCOMPATIBLE + return Competence.COMPATIBLE + \ No newline at end of file diff --git a/pymc3/step_methods/hmc/trajectory.py b/pymc3/step_methods/hmc/trajectory.py index a7c730a9c5..1654355c98 100644 --- a/pymc3/step_methods/hmc/trajectory.py +++ b/pymc3/step_methods/hmc/trajectory.py @@ -292,6 +292,197 @@ def _theano_single_leapfrog(H, q, p, q_grad, **theano_kwargs): return f +def get_theano_hamiltonian_manifold_functions(model_vars, shared, logpt, potential, **theano_kwargs): + """Construct theano functions for the Hamiltonian, energy, and leapfrog integrator when using Manifolds. + + Parameters + ---------- + model_vars : array of variables to be sampled + shared : theano tensors that are already shared + logpt : model log probability + potential : Hamiltonian potential + theano_kwargs : dictionary of keyword arguments to pass to theano functions + use_single_leapfrog : bool + if only 1 integration step is done at a time (as in NUTS), this + provides a ~2x speedup + integrator : str + Integration scheme to use. One of "leapfog", "two-stage", or + "three-stage". + + Returns + ------- + H : Hamiltonian namedtuple + energy_function : theano function computing energy at a point in phase space + leapfrog_integrator : theano function integrating the Hamiltonian from a point in phase space + theano_variables : dictionary of variables used in the computation graph which may be useful + """ + H, q, dlogp = _theano_hamiltonian_manifold(model_vars, shared, logpt, potential) + energy_function, p = _theano_energy_function_softabs(H, q, **theano_kwargs) + velocity_function = _theano_velocity_function_softabs(H, p, **theano_kwargs) + integrator = _theano_leapfrog_integrator_softabs(H, q, p, **theano_kwargs) + return H, energy_function, velocity_function, integrator, dlogp + + +def _theano_hamiltonian_manifold(model_vars, shared, logpt, potential): + """Creates a Hamiltonian with shared inputs. + + Parameters + ---------- + model_vars : array of variables to be sampled + shared : theano tensors that are already shared + logpt : model log probability + potential : hamiltonian potential + + Returns + ------- + Hamiltonian : namedtuple with log pdf, gradient of log pdf, and potential functions + q : Starting position variable. + """ + dlogp = gradient(logpt, model_vars) + (logp, dlogp), q = join_nonshared_inputs([logpt, dlogp], model_vars, shared) + dlogp_func = theano.function(inputs=[q], outputs=dlogp) + dlogp_func.trust_input = True + logp = CallableTensor(logp) + dlogp = CallableTensor(dlogp) + return Hamiltonian(logp, dlogp, potential), q, dlogp_func + + +def _theano_energy_function_softabs(H, q, **theano_kwargs): + """Creates a Hamiltonian with shared inputs. + + Parameters + ---------- + H : Hamiltonian namedtuple + q : theano variable, starting position + theano_kwargs : passed to theano.function + + Returns + ------- + energy_function : theano function that computes the energy at a point (p, q) in phase space + p : Starting momentum variable. + """ + p = tt.vector('p') + p.tag.test_value = q.tag.test_value + total_energy = H.pot.energy(p) - H.logp(q) + energy_function = theano.function(inputs=[q, p], outputs=total_energy, **theano_kwargs) + energy_function.trust_input = True + + return energy_function, p + + +def _theano_velocity_function_softabs(H, p, **theano_kwargs): + v = H.pot.velocity(p) + velocity_function = theano.function(inputs=[p], outputs=v, **theano_kwargs) + velocity_function.trust_input = True + return velocity_function + + +def energy_softabs(H, q, p): + """Compute the total energy for the Hamiltonian at a given position/momentum""" + return H.pot.energy(p) - H.logp(q) + + +def leapfrog_softabs(H, q, p, epsilon, n_steps): + """Leapfrog integrator. + + Estimates `p(t)` and `q(t)` at time :math:`t = n \cdot e`, by integrating the + Hamiltonian equations + + .. math:: + + \frac{dq_i}{dt} = \frac{\partial H}{\partial p_i} + + \frac{dp_i}{dt} = \frac{\partial H}{\partial q_i} + + with :math:`p(0) = p`, :math:`q(0) = q` + + Parameters + ---------- + H : Hamiltonian instance. + Tuple of `logp, dlogp, potential`. + q : Theano.tensor + initial position vector + p : Theano.tensor + initial momentum vector + epsilon : float, step size + n_steps : int, number of iterations + + Returns + ------- + position : Theano.tensor + position estimate at time :math:`n \cdot e`. + momentum : Theano.tensor + momentum estimate at time :math:`n \cdot e`. + """ + def full_update(p, q): + p = p + epsilon * H.dlogp(q) + q += epsilon * H.pot.velocity(p) + return p, q + # This first line can't be +=, possibly because of theano + p = p + 0.5 * epsilon * H.dlogp(q) # half momentum update + q += epsilon * H.pot.velocity(p) # full position update + if tt.gt(n_steps, 1): + (p_seq, q_seq), _ = theano.scan(full_update, outputs_info=[p, q], n_steps=n_steps - 1) + p, q = p_seq[-1], q_seq[-1] + p += 0.5 * epsilon * H.dlogp(q) # half momentum update + return q, p + + +def _theano_leapfrog_integrator_softabs(H, q, p, **theano_kwargs): + """Computes a theano function that computes one leapfrog step and the energy at the + end of the trajectory. + + Parameters + ---------- + H : Hamiltonian + q : theano.tensor + p : theano.tensor + theano_kwargs : passed to theano.function + + Returns + ------- + theano function which returns + q_new, p_new, energy_new + """ + epsilon = tt.scalar('epsilon') + epsilon.tag.test_value = 1. + + n_steps = tt.iscalar('n_steps') + n_steps.tag.test_value = 2 + + q_new, p_new = leapfrog_softabs(H, q, p, epsilon, n_steps) + energy_new = energy_softabs(H, q_new, p_new) + + f = theano.function([q, p, epsilon, n_steps], [q_new, p_new, energy_new], **theano_kwargs) + f.trust_input = True + return f + + +def quadratic_gradient(H_ij, p): + """ + pseudo-code of the gradient of the quadratic form p^T . H^-1 . p + To be used in the energy and velocity functions + """ + # Q, lambda_i = decompose(H_ij) + # D = diag(Q_t . p / (lambda_i . coth(alpha . lambda_i)) + # J = d(lambda_i . coth(alpha . lambda_i)) + # grad = - Trace(Q . D . J . D . Q_t . d(H)) + # return grad + return None + + +def log_gradient(H_ij): + """ + pseduo-code of the gradient of the log determinant. + To be used in the energy and velocity functions + """ + # Q, lambda_i = decompose(H_ij) + # J = d(lambda_i . coth(alpha . lambda_i)) + # R = diag(1 / lambda_i . coth(alpha . lambda_i) + # grad = Trace(Q . (R . J). Q_t . dH) + # return grad + return None + INTEGRATORS_SINGLE = { 'leapfrog': _theano_single_leapfrog, 'two-stage': _theano_single_twostage,