[WIP] Riemannian Manifold HMC

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)

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
+

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,