Refactor Hamiltonian methods into single class

Author: colin

pymc3/step_methods/__init__.py

@@ -1,6 +1,6 @@
 from .compound import CompoundStep
 
-from .hmc import HamiltonianMC
+from .hmc import HamiltonianMC, NUTS
 
 from .metropolis import Metropolis
 from .metropolis import BinaryMetropolis
@@ -15,5 +15,3 @@
 from .gibbs import ElemwiseCategorical
 
 from .slicer import Slice
-
-from .nuts import NUTS

pymc3/step_methods/arraystep.py

@@ -4,7 +4,6 @@
 from ..blocking import ArrayOrdering, DictToArrayBijection
 import numpy as np
 from numpy.random import uniform
-from numpy import log, isfinite
 from enum import IntEnum, unique
 
 __all__ = ['ArrayStep', 'ArrayStepShared', 'metrop_select', 'SamplerHist',
@@ -90,9 +89,9 @@ class ArrayStep(BlockedStep):
     ----------
     vars : list
         List of variables for sampler.
+    fs: list of logp theano functions
     allvars: Boolean (default False)
     blocked: Boolean (default True)
-    fs: logp theano function
     """
 
     def __init__(self, vars, fs, allvars=False, blocked=True):
@@ -114,10 +113,11 @@ def step(self, point):
 
 
 class ArrayStepShared(BlockedStep):
-    """Faster version of ArrayStep that requires the substep method that does not wrap the functions the step method uses.
+    """Faster version of ArrayStep that requires the substep method that does not wrap
+       the functions the step method uses.
 
-    Works by setting shared variables before using the step. This eliminates the mapping and unmapping overhead as well
-    as moving fewer variables around.
+    Works by setting shared variables before using the step. This eliminates the mapping
+    and unmapping overhead as well as moving fewer variables around.
     """
 
     def __init__(self, vars, shared, blocked=True):
@@ -144,14 +144,25 @@ def step(self, point):
 
 
 def metrop_select(mr, q, q0):
-    # Perform rejection/acceptance step for Metropolis class samplers
+    """Perform rejection/acceptance step for Metropolis class samplers.
 
+    Returns the new sample q if a uniform random number is less than the
+    metropolis acceptance rate (`mr`), and the old sample otherwise.
+
+    Parameters
+    ----------
+    mr : float, Metropolis acceptance rate
+    q : proposed sample
+    q0 : current sample
+
+    Returns
+    -------
+    q or q0
+    """
     # Compare acceptance ratio to uniform random number
-    if isfinite(mr) and log(uniform()) < mr:
-        # Accept proposed value
+    if np.isfinite(mr) and np.log(uniform()) < mr:
         return q
     else:
-        # Reject proposed value
         return q0
 
 

pymc3/step_methods/hmc.py

@@ -0,0 +1,2 @@
+from .hmc import HamiltonianMC
+from .nuts import NUTS

pymc3/step_methods/hmc/__init__.py

@@ -0,0 +1,57 @@
+from ..arraystep import ArrayStepShared
+from .trajectory import get_theano_hamiltonian_functions
+
+from pymc3.tuning import guess_scaling
+from pymc3.model import modelcontext, Point
+from .quadpotential import quad_potential
+from pymc3.theanof import inputvars, make_shared_replacements
+
+
+class BaseHMC(ArrayStepShared):
+    default_blocked = True
+
+    def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False,
+                 model=None, blocked=True, use_single_leapfrog=False, **theano_kwargs):
+        """
+        Parameters
+        ----------
+            vars : list of theano variables
+            scaling : array_like, ndim = {1,2}
+                Scaling for momentum distribution. 1d arrays interpreted matrix diagonal.
+            step_scale : float, default=0.25
+                Size of steps to take, automatically scaled down by 1/n**(1/4)
+            is_cov : bool, default=False
+                Treat scaling as a covariance matrix/vector if True, else treat it as a
+                precision matrix/vector
+            state
+                State object
+            model : pymc3 Model instance.  default=Context model
+            blocked: Boolean, default True
+            use_single_leapfrog: Boolean, will leapfrog steps take a single step at a time.
+                default False.
+            **theano_kwargs: passed to theano functions
+        """
+        model = modelcontext(model)
+
+        if vars is None:
+            vars = model.cont_vars
+        vars = inputvars(vars)
+
+        if scaling is None:
+            scaling = model.test_point
+
+        if isinstance(scaling, dict):
+            scaling = guess_scaling(Point(scaling, model=model), model=model, vars=vars)
+
+        n = scaling.shape[0]
+        self.step_size = step_scale / (n ** 0.25)
+        self.potential = quad_potential(scaling, is_cov, as_cov=False)
+
+        shared = make_shared_replacements(vars, model)
+        if theano_kwargs is None:
+            theano_kwargs = {}
+
+        self.H, self.compute_energy, self.leapfrog, self._vars = get_theano_hamiltonian_functions(
+            vars, shared, model.logpt, self.potential, use_single_leapfrog, **theano_kwargs)
+
+        super(BaseHMC, self).__init__(vars, shared, blocked=blocked)

pymc3/step_methods/hmc/base_hmc.py

@@ -0,0 +1,56 @@
+'''
+Created on Mar 7, 2011
+
+@author: johnsalvatier
+'''
+from ..arraystep import metrop_select, Competence
+from .base_hmc import BaseHMC
+from pymc3.vartypes import discrete_types
+import numpy as np
+
+
+__all__ = ['HamiltonianMC']
+
+# TODO:
+# add constraint handling via page 37 of Radford's
+# http://www.cs.utoronto.ca/~radford/ham-mcmc.abstract.html
+
+
+def unif(step_size, elow=.85, ehigh=1.15):
+    return np.random.uniform(elow, ehigh) * step_size
+
+
+class HamiltonianMC(BaseHMC):
+    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.
+            **kwargs : passed to BaseHMC
+        """
+        super(HamiltonianMC, self).__init__(vars, **kwargs)
+        self.path_length = path_length
+        self.step_rand = step_rand
+
+    def astep(self, q0):
+        e = np.array(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 = current_energy - initial_energy
+        return metrop_select(energy_change, q, q0)
+
+    @staticmethod
+    def competence(var):
+        if var.dtype in discrete_types:
+            return Competence.INCOMPATIBLE
+        return Competence.COMPATIBLE