Refactor Slice.astep (#1371)

Author: ColCarroll

pymc3/step_methods/slicer.py

@@ -1,11 +1,12 @@
 # Modified from original implementation by Dominik Wabersich (2013)
 
+import numpy as np
+import numpy.random as nr
+
 from .arraystep import ArrayStep, Competence
 from ..model import modelcontext
 from ..theanof import inputvars
 from ..vartypes import continuous_types
-from numpy import floor, abs, atleast_1d, empty, isfinite, sum, resize
-from numpy.random import standard_exponential, random, uniform
 
 __all__ = ['Slice']
 
@@ -28,66 +29,47 @@ class Slice(ArrayStep):
     """
     default_blocked = False
 
-    def __init__(self, vars=None, w=1, tune=True, model=None, **kwargs):
-
-        model = modelcontext(model)
+    def __init__(self, vars=None, w=1., tune=True, model=None, **kwargs):
+        self.model = modelcontext(model)
+        self.w = w
+        self.tune = tune
+        self.w_sum = 0
+        self.n_tunes = 0
 
         if vars is None:
-            vars = model.cont_vars
+            vars = self.model.cont_vars
         vars = inputvars(vars)
 
-        self.w = w
-        self.tune = tune
-        self.w_tune = []
-        self.model = model
-
-        super(Slice, self).__init__(vars, [model.fastlogp], **kwargs)
+        super(Slice, self).__init__(vars, [self.model.fastlogp], **kwargs)
 
     def astep(self, q0, logp):
-
-        q = q0.copy()
-        self.w = resize(self.w, len(q))
-
-        y = logp(q0) - standard_exponential()
+        self.w = np.resize(self.w, len(q0))
+        y = logp(q0) - nr.standard_exponential()
 
         # Stepping out procedure
-        ql = q0.copy()
-        ql -= uniform(0, self.w)
-        qr = q0.copy()
-        qr = ql + self.w
-
-        yl = logp(ql)
-        yr = logp(qr)
-
-        while((y < yl).all()):
-            ql -= self.w
-            yl = logp(ql)
+        q_left = q0 - nr.uniform(0, self.w)
+        q_right = q_left + self.w
 
-        while((y < yr).all()):
-            qr += self.w
-            yr = logp(qr)
+        while (y < logp(q_left)).all():
+            q_left -= self.w
 
-        q_next = q0.copy()
-        while True:
+        while (y < logp(q_right)).all():
+            q_right += self.w
 
+        q = nr.uniform(q_left, q_right, size=q_left.size)  # new variable to avoid copies
+        while logp(q) <= y:
             # Sample uniformly from slice
-            qi = uniform(ql, qr, size=ql.size)
-
-            yi = logp(qi)
-
-            if yi > y:
-                q = qi
-                break
-            elif (qi > q).all():
-                qr = qi
-            elif (qi < q).all():
-                ql = qi
+            if (q > q0).all():
+                q_right = q
+            elif (q < q0).all():
+                q_left = q
+            q = nr.uniform(q_left, q_right, size=q_left.size)
 
         if self.tune:
             # Tune sampler parameters
-            self.w_tune.append(abs(q0 - q))
-            self.w = 2 * sum(self.w_tune, 0) / len(self.w_tune)
-
+            self.w_sum += np.abs(q0 - q)
+            self.n_tunes += 1.
+            self.w = 2. * self.w_sum / self.n_tunes
         return q
 
     @staticmethod