diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index 23c0a8d67f..517baaec7f 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/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) diff --git a/pymc3/step_methods/hmc/integration.py b/pymc3/step_methods/hmc/integration.py index 04605071b1..7c40c2163a 100644 --- a/pymc3/step_methods/hmc/integration.py +++ b/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) \ No newline at end of file diff --git a/pymc3/step_methods/hmc/rmhmc.py b/pymc3/step_methods/hmc/rmhmc.py new file mode 100644 index 0000000000..9106e8c60f --- /dev/null +++ b/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