[WIP] RMHMC

pymc3/step_methods/hmc/base_hmc.py

@@ -12,7 +12,7 @@ class BaseHMC(arraystep.GradientSharedStep):
 
     def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False,
                  model=None, blocked=True, potential=None,
-                 integrator="leapfrog", dtype=None, **theano_kwargs):
+                 integrator="leapfrog", dtype=None, manifold=False, **theano_kwargs):
         """Superclass to implement Hamiltonian/hybrid monte carlo
 
         Parameters
@@ -62,5 +62,9 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False,
         else:
             self.potential = quad_potential(scaling, is_cov)
 
-        self.integrator = integration.CpuLeapfrogIntegrator(
-            size, self.potential, self._logp_dlogp_func)
+        if manifold:
+            self.integrator = integration.CpuLeapfrogIntegratorSoftabs(
+                size, self.potential, self._logp_dlogp_func)
+        else:
+            self.integrator = integration.CpuLeapfrogIntegrator(
+                size, self.potential, self._logp_dlogp_func)

pymc3/step_methods/hmc/integration.py

@@ -66,3 +66,64 @@ def step(self, epsilon, state, out=None):
             return
         else:
             return State(q_new, p_new, v_new, q_new_grad, energy)
+
+
+class CpuLeapfrogIntegratorSoftabs(object):
+    def __init__(self, ndim, potential, logp_dlogp_func):
+        self._ndim = ndim
+        self._potential = potential
+        self._logp_dlogp_func = logp_dlogp_func
+        self._dtype = self._logp_dlogp_func.dtype
+        if self._potential.dtype != self._dtype:
+            raise ValueError("dtypes of potential and logp function "
+                             "don't match."
+                             % (self._potential.dtype, self._dtype))
+
+    def compute_state(self, q, p):
+        if q.dtype != self._dtype or p.dtype != self._dtype:
+            raise ValueError('Invalid dtype. Must be %s' % self._dtype)
+        logp, dlogp = self._logp_dlogp_func(q)
+        v = self._potential.velocity(p)
+        kinetic = self._potential.energy(p, velocity=v)
+        energy = kinetic - logp
+        return State(q, p, v, dlogp, energy)
+
+    def step(self, epsilon, state, out=None):
+        pot = self._potential
+        axpy = linalg.blas.get_blas_funcs('axpy', dtype=self._dtype)
+
+        q, p, v, q_grad, energy = state
+        if out is None:
+            q_new = q.copy()
+            p_new = p.copy()
+            v_new = np.empty_like(q)
+            q_new_grad = np.empty_like(q)
+        else:
+            q_new, p_new, v_new, q_new_grad, energy = out
+            q_new[:] = q
+            p_new[:] = p
+
+        dt = 0.5 * epsilon
+
+        # p is already stored in p_new
+        # p_new = p + dt * q_grad
+        axpy(q_grad, p_new, a=dt)
+
+        pot.velocity(p_new, out=v_new)
+        # q is already stored in q_new
+        # q_new = q + epsilon * v_new
+        axpy(v_new, q_new, a=epsilon)
+
+        logp = self._logp_dlogp_func(q_new, q_new_grad)
+
+        # p_new = p_new + dt * q_new_grad
+        axpy(q_new_grad, p_new, a=dt)
+
+        kinetic = pot.velocity_energy(p_new, v_new)
+        energy = kinetic - logp
+
+        if out is not None:
+            out.energy = energy
+            return
+        else:
+            return State(q_new, p_new, v_new, q_new_grad, energy)

pymc3/step_methods/hmc/rmhmc.py

@@ -0,0 +1,72 @@
+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 = 'rmmc'
+    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 = int(self.path_length / e)
+
+        p0 = self.potential.random()
+        start = self.integrator.compute_state(q0, p0)
+
+        if not np.isfinite(start.energy):
+            raise ValueError('Bad initial energy: %s. The model '
+                             'might be misspecified.' % start.energy)
+
+        state = start
+        for _ in range(n_steps):
+            state = self.integrator.step(e, state)
+
+        energy_change = start.energy - state.energy
+        return metrop_select(energy_change, state.q, start.q)[0]
+
+    @staticmethod
+    def competence(var):
+        if var.dtype in discrete_types:
+            return Competence.INCOMPATIBLE
+        return Competence.COMPATIBLE