diff --git a/pymc3/model.py b/pymc3/model.py index ed922e63f3..d3f9dd96bf 100644 --- a/pymc3/model.py +++ b/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): diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index d5186eeb7c..e057be3ac0 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/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: diff --git a/pymc3/step_methods/hmc/quadpotential.py b/pymc3/step_methods/hmc/quadpotential.py index 835e44174c..554b4d2028 100644 --- a/pymc3/step_methods/hmc/quadpotential.py +++ b/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") diff --git a/pymc3/tests/test_quadpotential.py b/pymc3/tests/test_quadpotential.py index 9a7a5da052..53b0b0b1f8 100644 --- a/pymc3/tests/test_quadpotential.py +++ b/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