Allow user defined potentials for hmc

pymc3/model.py

@@ -412,6 +412,10 @@ def bijection(self):
     def dict_to_array(self):
         return self.bijection.map
 
+    @property
+    def ndim(self):
+        return self.dict_to_array(self.test_point).shape[0]
+
     @property
     @memoize
     def logp_array(self):

pymc3/step_methods/hmc/base_hmc.py

@@ -11,7 +11,8 @@ 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):
+                 model=None, blocked=True, use_single_leapfrog=False,
+                 potential=None, **theano_kwargs):
         """Superclass to implement Hamiltonian/hybrid monte carlo
 
         Parameters
@@ -30,6 +31,9 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False,
         blocked: Boolean, default True
         use_single_leapfrog: Boolean, will leapfrog steps take a single step at a time.
             default False.
+        potential : Potential, optional
+            An object that represents the Hamiltonian with methods `velocity`,
+            `energy`, and `random` methods.
         **theano_kwargs: passed to theano functions
         """
         model = modelcontext(model)
@@ -38,15 +42,20 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False,
             vars = model.cont_vars
         vars = inputvars(vars)
 
-        if scaling is None:
+        if scaling is None and potential 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)
+        if scaling is not None and potential is not None:
+            raise ValueError("Can not specify both potential and scaling.")
+
+        self.step_size = step_scale / (model.ndim ** 0.25)
+        if potential is not None:
+            self.potential = potential
+        else:
+            self.potential = quad_potential(scaling, is_cov, as_cov=False)
 
         shared = make_shared_replacements(vars, model)
         if theano_kwargs is None:

pymc3/step_methods/hmc/quadpotential.py

@@ -18,7 +18,7 @@ def quad_potential(C, is_cov, as_cov):
     ----------
     C : arraylike, 0 <= ndim <= 2
         scaling matrix for the potential
-        vector treated as diagonal matrix
+        vector treated as diagonal matrix.
     is_cov : Boolean
         whether C is provided as a covariance matrix or hessian
     as_cov : Boolean
@@ -29,7 +29,6 @@ def quad_potential(C, is_cov, as_cov):
     -------
     q : Quadpotential
     """
-
     if issparse(C):
         if not chol_available:
             raise ImportError("Sparse mass matrices require scikits.sparse")

pymc3/tests/test_quadpotential.py

@@ -5,6 +5,7 @@
 import theano
 
 from pymc3.step_methods.hmc import quadpotential
+import pymc3
 
 from nose.tools import raises
 from nose.plugins.skip import SkipTest
@@ -147,3 +148,23 @@ def test_random_dense():
         for pot in pots:
             cov_ = np.cov(np.array([pot.random() for _ in range(1000)]).T)
             assert np.allclose(cov_, inv, atol=0.1)
+
+
+def test_user_potential():
+    model = pymc3.Model()
+    with model:
+        a = pymc3.Normal("a", mu=0, sd=1)
+
+    # Work around missing nonlocal in python2
+    called = []
+
+    class Potential(quadpotential.ElemWiseQuadPotential):
+        def energy(self, x):
+            called.append(1)
+            return super(Potential, self).energy(x)
+
+    pot = Potential([1])
+    with model:
+        step = pymc3.NUTS(potential=pot)
+        pymc3.sample(10, init=None, step=step)
+    assert called