Clean up HMC code (#3447)

* Clean up HMC code

* Fix some tests

* Add back target_accept attribute

Author: ColCarroll

pymc3/model.py

@@ -492,7 +492,7 @@ def __call__(self, array, grad_out=None, extra_vars=None):
         if grad_out is None:
             return logp, dlogp
         else:
-            out[...] = dlogp
+            np.copyto(out, dlogp)
             return logp
 
     @property
@@ -737,7 +737,7 @@ def logpt(self):
     def logp_nojact(self):
         """Theano scalar of log-probability of the model but without the jacobian
         if transformed Random Variable is presented.
-        Note that If there is no transformed variable in the model, logp_nojact 
+        Note that If there is no transformed variable in the model, logp_nojact
         will be the same as logpt as there is no need for Jacobian correction.
         """
         with self:

pymc3/step_methods/hmc/base_hmc.py

@@ -19,7 +19,6 @@
 
 DivergenceInfo = namedtuple("DivergenceInfo", "message, exec_info, state")
 
-
 class BaseHMC(arraystep.GradientSharedStep):
     """Superclass to implement Hamiltonian/hybrid monte carlo."""
 
@@ -34,7 +33,6 @@ def __init__(
         model=None,
         blocked=True,
         potential=None,
-        integrator="leapfrog",
         dtype=None,
         Emax=1000,
         target_accept=0.8,
@@ -79,11 +77,10 @@ def __init__(
         size = self._logp_dlogp_func.size
 
         self.step_size = step_scale / (size ** 0.25)
-        self.target_accept = target_accept
         self.step_adapt = step_sizes.DualAverageAdaptation(
             self.step_size, target_accept, gamma, k, t0
         )
-
+        self.target_accept = target_accept
         self.tune = True
 
         if scaling is None and potential is None:

pymc3/step_methods/hmc/integration.py

@@ -32,7 +32,7 @@ def compute_state(self, q, p):
         energy = kinetic - logp
         return State(q, p, v, dlogp, energy, logp)
 
-    def step(self, epsilon, state, out=None):
+    def step(self, epsilon, state):
         """Leapfrog integrator step.
 
         Half a momentum update, full position update, half momentum update.
@@ -51,7 +51,7 @@ def step(self, epsilon, state, out=None):
         None if `out` is provided, else a State namedtuple
         """
         try:
-            return self._step(epsilon, state, out=None)
+            return self._step(epsilon, state)
         except linalg.LinAlgError as err:
             msg = "LinAlgError during leapfrog step."
             raise IntegrationError(msg)
@@ -64,26 +64,20 @@ def step(self, epsilon, state, out=None):
             else:
                 raise
 
-    def _step(self, epsilon, state, out=None):
-        pot = self._potential
+    def _step(self, epsilon, state):
         axpy = linalg.blas.get_blas_funcs('axpy', dtype=self._dtype)
+        pot = self._potential
 
-        q, p, v, q_grad, energy, logp = 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
+        q_new = state.q.copy()
+        p_new = state.p.copy()
+        v_new = np.empty_like(q_new)
+        q_new_grad = np.empty_like(q_new)
 
         dt = 0.5 * epsilon
 
         # p is already stored in p_new
         # p_new = p + dt * q_grad
-        axpy(q_grad, p_new, a=dt)
+        axpy(state.q_grad, p_new, a=dt)
 
         pot.velocity(p_new, out=v_new)
         # q is already stored in q_new
@@ -98,8 +92,4 @@ def _step(self, epsilon, state, out=None):
         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, logp)
+        return State(q_new, p_new, v_new, q_new_grad, energy, logp)

pymc3/tests/test_distributions_timeseries.py

@@ -5,6 +5,10 @@
 from ..theanof import floatX
 
 import numpy as np
+import pytest
+
+pytestmark = pytest.mark.usefixtures('seeded_test')
+
 
 def test_AR():
     # AR1

pymc3/tests/test_posteriors.py

@@ -34,19 +34,11 @@ class TestSliceUniform(sf.SliceFixture, sf.UniformFixture):
 
 
 class TestNUTSUniform2(TestNUTSUniform):
-    step_args = {'target_accept': 0.95, 'integrator': 'two-stage'}
+    step_args = {'target_accept': 0.95}
 
 
 class TestNUTSUniform3(TestNUTSUniform):
-    step_args = {'target_accept': 0.80, 'integrator': 'two-stage'}
-
-
-class TestNUTSUniform4(TestNUTSUniform):
-    step_args = {'target_accept': 0.95, 'integrator': 'three-stage'}
-
-
-class TestNUTSUniform5(TestNUTSUniform):
-    step_args = {'target_accept': 0.80, 'integrator': 'three-stage'}
+    step_args = {'target_accept': 0.80}
 
 
 class TestNUTSNormal(sf.NutsFixture, sf.NormalFixture):