Use sum of momentum for nuts termination

Author: aseyboldt

pymc3/step_methods/hmc/base_hmc.py

@@ -59,7 +59,7 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False,
         if theano_kwargs is None:
             theano_kwargs = {}
 
-        self.H, self.compute_energy, self.leapfrog, self.dlogp = get_theano_hamiltonian_functions(
+        self.H, self.compute_energy, self.compute_velocity, self.leapfrog, self.dlogp = get_theano_hamiltonian_functions(
             vars, shared, model.logpt, self.potential, use_single_leapfrog, integrator, **theano_kwargs)
 
         super(BaseHMC, self).__init__(vars, shared, blocked=blocked)

pymc3/step_methods/hmc/nuts.py

@@ -79,7 +79,8 @@ class NUTS(BaseHMC):
     }]
 
     def __init__(self, vars=None, Emax=1000, target_accept=0.8,
-                 gamma=0.05, k=0.75, t0=10, adapt_step_size=True, **kwargs):
+                 gamma=0.05, k=0.75, t0=10, adapt_step_size=True,
+                 max_treedepth=10, **kwargs):
         """
         Parameters
         ----------
@@ -124,11 +125,13 @@ def __init__(self, vars=None, Emax=1000, target_accept=0.8,
         self.log_step_size_bar = 0
         self.m = 1
         self.adapt_step_size = adapt_step_size
+        self.max_treedepth = max_treedepth
 
         self.tune = True
 
     def astep(self, q0):
         p0 = self.potential.random()
+        v0 = self.compute_velocity(p0)
         start_energy = self.compute_energy(q0, p0)
 
         if not self.adapt_step_size:
@@ -138,10 +141,10 @@ def astep(self, q0):
         else:
             step_size = np.exp(self.log_step_size_bar)
 
-        start = Edge(q0, p0, self.dlogp(q0), start_energy)
-        tree = Tree(self.leapfrog, start, step_size, self.Emax)
+        start = Edge(q0, p0, v0, self.dlogp(q0), start_energy)
+        tree = Tree(len(p0), self.leapfrog, start, step_size, self.Emax)
 
-        while True:
+        for _ in range(self.max_treedepth):
             direction = logbern(np.log(0.5)) * 2 - 1
             diverging, turning = tree.extend(direction)
             q = tree.proposal.q
@@ -179,17 +182,17 @@ def competence(var):
 
 
 # A node in the NUTS tree that is at the far right or left of the tree
-Edge = namedtuple("Edge", 'q, p, q_grad, energy')
+Edge = namedtuple("Edge", 'q, p, v, q_grad, energy')
 
 # A proposal for the next position
 Proposal = namedtuple("Proposal", "q, energy, p_accept")
 
 # A subtree of the binary tree build by nuts.
-Subtree = namedtuple("Subtree", "left, right, proposal, depth, log_size, accept_sum, n_proposals")
+Subtree = namedtuple("Subtree", "left, right, p_sum, proposal, log_size, accept_sum, n_proposals")
 
 
 class Tree(object):
-    def __init__(self, leapfrog, start, step_size, Emax):
+    def __init__(self, ndim, leapfrog, start, step_size, Emax):
         """Binary tree from the NUTS algorithm.
 
         Parameters
@@ -204,6 +207,7 @@ def __init__(self, leapfrog, start, step_size, Emax):
             The maximum energy change to accept before aborting the
             transition as diverging.
         """
+        self.ndim = ndim
         self.leapfrog = leapfrog
         self.start = start
         self.step_size = step_size
@@ -214,9 +218,9 @@ def __init__(self, leapfrog, start, step_size, Emax):
         self.proposal = Proposal(start.q, start.energy, 1.0)
         self.depth = 0
         self.log_size = 0
-        # TODO Why not a global accept sum and n_proposals?
-        #self.accept_sum = 0
-        #self.n_proposals = 0
+        self.accept_sum = 0
+        self.n_proposals = 0
+        self.p_sum = start.p.copy()
         self.max_energy_change = 0
 
     def extend(self, direction):
@@ -237,63 +241,75 @@ def extend(self, direction):
             self.right = tree.right
         else:
             tree, diverging, turning = self._build_subtree(
-                self.left, self.depth, floatX(np.asarray(- self.step_size)))
+                self.left, self.depth, floatX(np.asarray(-self.step_size)))
             self.left = tree.right
 
-        ok = not (diverging or turning)
-        if ok and logbern(tree.log_size - self.log_size):
+        self.depth += 1
+        self.accept_sum += tree.accept_sum
+        self.n_proposals += tree.n_proposals
+
+        if diverging or turning:
+            return diverging, turning
+
+        size1, size2 = self.log_size, tree.log_size
+        if logbern(size2 - size1):
             self.proposal = tree.proposal
 
-        self.depth += 1
         self.log_size = np.logaddexp(self.log_size, tree.log_size)
-        # TODO why not +=
-        #self.accept_sum += tree.accept_sum
-        self.accept_sum = tree.accept_sum
-        #self.n_proposals += tree.n_proposals
-        self.n_proposals = tree.n_proposals
+        self.p_sum[:] += tree.p_sum
 
         left, right = self.left, self.right
-        span = right.q - left.q
-        turning = turning or (span.dot(left.p) < 0) or (span.dot(right.p) < 0)
+        p_sum = self.p_sum
+        turning = (p_sum.dot(left.v) <= 0) or (p_sum.dot(right.v) <= 0)
+
         return diverging, turning
 
     def _build_subtree(self, left, depth, epsilon):
         if depth == 0:
             right = self.leapfrog(left.q, left.p, left.q_grad, epsilon)
             right = Edge(*right)
             energy_change = right.energy - self.start_energy
+            if np.isnan(energy_change):
+                energy_change = np.inf
+
             if np.abs(energy_change) > np.abs(self.max_energy_change):
                 self.max_energy_change = energy_change
             p_accept = min(1, np.exp(-energy_change))
 
             log_size = -energy_change
-            diverging = not np.isfinite(energy_change)
-            diverging = diverging or (np.abs(energy_change) > self.Emax)
+            diverging = energy_change > self.Emax
 
             proposal = Proposal(right.q, right.energy, p_accept)
-            tree = Subtree(right, right, proposal, 1, log_size, p_accept, 1)
+            tree = Subtree(right, right, right.p, proposal, log_size, p_accept, 1)
             return tree, diverging, False
 
         tree1, diverging, turning = self._build_subtree(left, depth - 1, epsilon)
         if diverging or turning:
             return tree1, diverging, turning
 
         tree2, diverging, turning = self._build_subtree(tree1.right, depth - 1, epsilon)
+        ok = not (diverging or turning)
 
-        log_size = np.logaddexp(tree1.log_size, tree2.log_size)
         accept_sum = tree1.accept_sum + tree2.accept_sum
         n_proposals = tree1.n_proposals + tree2.n_proposals
 
         left, right = tree1.left, tree2.right
-        span = np.sign(epsilon) * (right.q - left.q)
-        turning = turning or (span.dot(left.p) < 0) or (span.dot(right.p) < 0)
 
-        if np.isfinite(tree2.log_size) and logbern(tree2.log_size - log_size):
-            proposal = tree2.proposal
+        if ok:
+            p_sum = tree1.p_sum + tree2.p_sum
+            turning = (p_sum.dot(left.v) <= 0) or (p_sum.dot(right.v) <= 0)
+
+            log_size = np.logaddexp(tree1.log_size, tree2.log_size)
+            if logbern(tree2.log_size - log_size):
+                proposal = tree2.proposal
+            else:
+                proposal = tree1.proposal
         else:
+            p_sum = tree1.p_sum
+            log_size = tree1.log_size
             proposal = tree1.proposal
 
-        tree = Subtree(left, right, proposal, depth, log_size, accept_sum, n_proposals)
+        tree = Subtree(left, right, p_sum, proposal, log_size, accept_sum, n_proposals)
         return tree, diverging, turning
 
     def stats(self):

pymc3/step_methods/hmc/quadpotential.py

@@ -1,6 +1,7 @@
 from numpy import dot
 from numpy.random import normal
 import scipy.linalg
+import theano.tensor as tt
 from theano.tensor import slinalg
 from scipy.sparse import issparse
 
@@ -123,7 +124,7 @@ def __init__(self, A):
         self.L = scipy.linalg.cholesky(A, lower=True)
 
     def velocity(self, x):
-        return x.T.dot(self.A.T)
+        return tt.dot(self.A, x)
 
     def random(self):
         n = floatX(normal(size=self.L.shape[0]))

pymc3/step_methods/hmc/trajectory.py

@@ -57,6 +57,13 @@ def _theano_energy_function(H, q, **theano_kwargs):
     return energy_function, p
 
 
+def _theano_velocity_function(H, p, **theano_kwargs):
+    v = H.pot.velocity(p)
+    velocity_function = theano.function(inputs=[p], outputs=v, **theano_kwargs)
+    velocity_function.trust_input = True
+    return velocity_function
+
+
 def _theano_leapfrog_integrator(H, q, p, **theano_kwargs):
     """Computes a theano function that computes one leapfrog step and the energy at the
     end of the trajectory.
@@ -115,6 +122,7 @@ def get_theano_hamiltonian_functions(model_vars, shared, logpt, potential,
     """
     H, q, dlogp = _theano_hamiltonian(model_vars, shared, logpt, potential)
     energy_function, p = _theano_energy_function(H, q, **theano_kwargs)
+    velocity_function = _theano_velocity_function(H, p, **theano_kwargs)
     if use_single_leapfrog:
         try:
             _theano_integrator = INTEGRATORS_SINGLE[integrator]
@@ -125,7 +133,7 @@ def get_theano_hamiltonian_functions(model_vars, shared, logpt, potential,
         if integrator != "leapfrog":
             raise ValueError("Only leapfrog is supported")
         integrator = _theano_leapfrog_integrator(H, q, p, **theano_kwargs)
-    return H, energy_function, integrator, dlogp
+    return H, energy_function, velocity_function, integrator, dlogp
 
 
 def energy(H, q, p):
@@ -214,11 +222,12 @@ def _theano_single_threestage(H, q, p, q_grad, **theano_kwargs):
     q_e = q_1be + floatX(b) * epsilon * H.pot.velocity(p_1ae)
     grad_e = H.dlogp(q_e)
     p_e = p_1ae + floatX(a) * epsilon * grad_e
+    v_e = H.pot.velocity(p_e)
 
     new_energy = energy(H, q_e, p_e)
 
     f = theano.function(inputs=[q, p, q_grad, epsilon],
-                        outputs=[q_e, p_e, grad_e, new_energy],
+                        outputs=[q_e, p_e, v_e, grad_e, new_energy],
                         **theano_kwargs)
     f.trust_input = True
     return f
@@ -250,10 +259,11 @@ def _theano_single_twostage(H, q, p, q_grad, **theano_kwargs):
     q_e = q_e2 + epsilon / 2 * H.pot.velocity(p_1ae)
     grad_e = H.dlogp(q_e)
     p_e = p_1ae + a * epsilon * grad_e
+    v_e = H.pot.velocity(p_e)
 
     new_energy = energy(H, q_e, p_e)
     f = theano.function(inputs=[q, p, q_grad, epsilon],
-                        outputs=[q_e, p_e, grad_e, new_energy],
+                        outputs=[q_e, p_e, v_e, grad_e, new_energy],
                         **theano_kwargs)
     f.trust_input = True
     return f
@@ -273,15 +283,47 @@ def _theano_single_leapfrog(H, q, p, q_grad, **theano_kwargs):
     q_new_grad = H.dlogp(q_new)
     p_new += 0.5 * epsilon * q_new_grad  # half momentum update
     energy_new = energy(H, q_new, p_new)
+    v_new = H.pot.velocity(p_new)
 
     f = theano.function(inputs=[q, p, q_grad, epsilon],
-                        outputs=[q_new, p_new, q_new_grad, energy_new], **theano_kwargs)
+                        outputs=[q_new, p_new, v_new, q_new_grad, energy_new],
+                        **theano_kwargs)
+    f.trust_input = True
+    return f
+
+
+def _theano_single_leapfrog3(H, q, p, q_grad, **theano_kwargs):
+    """Do three leapfrog steps."""
+    step_size = tt.scalar('epsilon')
+    step_size.tag.test_value = 1.
+
+    epsilon = step_size / 3
+
+    p_new = p + 0.5 * epsilon * q_grad  # half momentum update
+    q_new = q + epsilon * H.pot.velocity(p_new)  # full position update
+
+    p_new = p_new + epsilon * H.dlogp(q_new)
+    q_new = q_new + epsilon * H.pot.velocity(p_new)
+
+    p_new = p_new + epsilon * H.dlogp(q_new)
+    q_new = q_new + epsilon * H.pot.velocity(p_new)
+
+    q_new_grad = H.dlogp(q_new)
+    p_new = p_new + 0.5 * epsilon * q_new_grad
+
+    energy_new = energy(H, q_new, p_new)
+    v_new = H.pot.velocity(p_new)
+
+    f = theano.function(inputs=[q, p, q_grad, step_size],
+                        outputs=[q_new, p_new, v_new, q_new_grad, energy_new],
+                        **theano_kwargs)
     f.trust_input = True
     return f
 
 
 INTEGRATORS_SINGLE = {
     'leapfrog': _theano_single_leapfrog,
     'two-stage': _theano_single_twostage,
-    'three-stage': _theano_single_threestage
+    'three-stage': _theano_single_threestage,
+    'leapfrog3': _theano_single_leapfrog3,
 }

pymc3/tests/sampler_fixtures.py

@@ -57,10 +57,6 @@ def make_model(cls):
             a = pm.Uniform("a", lower=-1, upper=1)
         return model
 
-    def test_interval(self):
-        a = self.samples['a']
-        npt.assert_almost_equal(((a > 0.1) & (a < 0.5)).mean(), 0.2, 2)
-
 
 class NormalFixture(KnownMean, KnownVariance, KnownCDF):
     means = {'a': 2 * np.ones(10)}
@@ -86,9 +82,8 @@ def make_model(cls):
         return model
 
 
-class StudentTFixture(KnownMean, KnownVariance, KnownCDF):
+class StudentTFixture(KnownMean, KnownCDF):
     means = {'a': 0}
-    variances = {'a': 3}
     cdfs = {'a': stats.t(df=3).cdf}
     ks_thin = 10
 

pymc3/tests/test_hmc.py

@@ -30,7 +30,7 @@ def test_leapfrog_reversible_single():
     n = 3
     start, model, _ = models.non_normal(n)
 
-    integrators = ['leapfrog', 'two-stage', 'three-stage']
+    integrators = ['leapfrog', 'two-stage', 'three-stage', 'leapfrog3']
     steps = [BaseHMC(vars=model.vars, model=model, integrator=method, use_single_leapfrog=True)
              for method in integrators]
     for method, step in zip(integrators, steps):
@@ -46,10 +46,10 @@ def test_leapfrog_reversible_single():
 
                 energy = step.compute_energy(q, p)
                 for _ in range(n_steps):
-                    q, p, dlogp, _ = step.leapfrog(q, p, dlogp, np.array(epsilon))
+                    q, p, v, dlogp, _ = step.leapfrog(q, p, dlogp, np.array(epsilon))
                 p = -p
                 for _ in range(n_steps):
-                    q, p, dlogp, _ = step.leapfrog(q, p, dlogp, np.array(epsilon))
+                    q, p, v, dlogp, _ = step.leapfrog(q, p, dlogp, np.array(epsilon))
 
                 close_to(q, q0, 1e-8, str(('q', method, n_steps, epsilon)))
                 close_to(-p, p0, 1e-8, str(('p', method, n_steps, epsilon)))

pymc3/tests/test_posteriors.py

@@ -62,13 +62,12 @@ class NUTSNormal(sf.NutsFixture, sf.NormalFixture):
 
 
 class NUTSBetaBinomial(sf.NutsFixture, sf.BetaBinomialFixture):
-    n_samples = 10000
+    n_samples = 2000
+    ks_thin = 5
     tune = 1000
     burn = 1000
     chains = 2
-    min_n_eff = 2000
-    rtol = 0.1
-    atol = 0.05
+    min_n_eff = 400
 
 
 @attr('extra')