Merge pull request #587 from pymc-devs/blocked_sampling

Implement non-blocked sampling for array based samplers.

Author: twiecki

pymc/examples/ARM12_6.py

@@ -63,6 +63,6 @@ def run(n=3000):
         step = HamiltonianMC(model.vars, h)
 
         trace = sample(n, step, start)
-        
+
 if __name__ == '__main__':
     run()

pymc/examples/ARM12_6uranium.py

@@ -69,4 +69,4 @@ def run(n=3000):
 
         trace = sample(n, step, start)
 if __name__ == '__main__':
-    run()      
+    run()

pymc/examples/arbitrary_stochastic.py

@@ -19,9 +19,8 @@ def run (n=3000):
 
         start = model.test_point
         h = find_hessian(start)
-        step = Metropolis(model.vars, h)
+        step = Metropolis(model.vars, h, blocked=True)
         trace = sample(n, step, start)
 
 if __name__ == "__main__":
     run()
-

pymc/examples/banana.py

@@ -63,5 +63,3 @@ def run(n = 3000):
         pl.figure()
         extent = np.min(xs), np.max(xs), np.min(ys), np.max(ys)
         pl.imshow(post, extent=extent)
-
-

pymc/examples/disaster_model.py

@@ -61,5 +61,3 @@ def run(n=1000):
 
 if __name__ == '__main__':
     run()
-
-

pymc/examples/factor_potential.py

@@ -13,7 +13,6 @@ def run(n = 3000):
         n = 50
     with model:
         trace = sample(n, step, start)
+
 if __name__ == '__main__':
     run()
-
-

pymc/examples/glm_linear.py

@@ -31,7 +31,7 @@ def run(n=2000):
     import matplotlib.pyplot as plt
 
     with model:
-        trace = sample(n, Slice(model.vars))
+        trace = sample(n, Slice())
 
     plt.plot(x, y, 'x')
     glm.plot_posterior_predictive(trace)

pymc/examples/glm_robust.py

@@ -37,7 +37,7 @@ def run(n=2000):
     import matplotlib.pyplot as plt
 
     with model:
-        trace = sample(n, Slice(model.vars))
+        trace = sample(n, Slice())
 
     plt.plot(x, y, 'x')
     glm.plot_posterior_predictive(trace)

pymc/examples/lasso_block_update.py

@@ -43,7 +43,7 @@
 # <codecell>
 
 with model:
-    step1 = Metropolis([m1, m2])
+    step1 = Metropolis([m1, m2], blocked=True)
 
     step2 = Metropolis([s], proposal_dist=LaplaceProposal)
 
@@ -67,5 +67,3 @@ def run(n=5000):
 # <codecell>
 if __name__ == '__main__':
     run()
-
-

pymc/examples/latent_occupancy.py

@@ -85,6 +85,6 @@ def run(n=5000):
         step2 = BinaryMetropolis([z])
 
         trace = sample(n, [step1, step2], start)
-        
+
 if __name__ == '__main__':
     run()

pymc/examples/stochastic_volatility.py

@@ -112,9 +112,9 @@ def hessian(point, nusd):
 
 with model:
     step = NUTS(model.vars, hessian(start, 6))
-    
-    
-    
+
+
+
 def run(n=2000):
     if n == "short":
         n = 50
@@ -136,7 +136,7 @@ def run(n=2000):
 
     # figsize(12,6)
     traceplot(trace, model.vars[:-1])
-    
+
 if __name__ == '__main__':
     run()
 

pymc/step_methods/arraystep.py

@@ -1,4 +1,6 @@
 from ..core import *
+from .compound import CompoundStep
+
 import numpy as np
 from numpy.random import uniform
 from numpy import log, isfinite
@@ -9,10 +11,44 @@
 
 
 class ArrayStep(object):
-    def __init__(self, vars, fs, allvars=False):
+    def __new__(cls, *args, **kwargs):
+        blocked = kwargs.get('blocked')
+        if blocked is None:
+            # Try to look up default value from class
+            blocked = getattr(cls, 'default_blocked', True)
+            kwargs['blocked'] = blocked
+
+        model = modelcontext(kwargs.get('model'))
+
+        # vars can either be first arg or a kwarg
+        if 'vars' not in kwargs and len(args) >= 1:
+            vars = args[0]
+            args = args[1:]
+        elif 'vars' in kwargs:
+            vars = kwargs.pop('vars')
+        else: # Assume all model variables
+            vars = model.vars
+
+        if not blocked and len(vars) > 1:
+            # In this case we create a separate sampler for each var
+            # and append them to a CompoundStep
+            steps = []
+            for var in vars:
+                step = super(ArrayStep, cls).__new__(cls)
+                # If we don't return the instance we have to manually
+                # call __init__
+                step.__init__([var], *args, **kwargs)
+                steps.append(step)
+
+            return CompoundStep(steps)
+        else:
+            return super(ArrayStep, cls).__new__(cls)
+
+    def __init__(self, vars, fs, allvars=False, blocked=True):
         self.ordering = ArrayOrdering(vars)
         self.fs = fs
         self.allvars = allvars
+        self.blocked = blocked
 
     def step(self, point):
         bij = DictToArrayBijection(self.ordering, point)

pymc/step_methods/hmc.py

@@ -26,7 +26,9 @@ def unif(step_size, elow=.85, ehigh=1.15):
 
 
 class HamiltonianMC(ArrayStep):
-    def __init__(self, vars=None, scaling=None, step_scale=.25, path_length=2., is_cov=False, step_rand=unif, state=None, model=None):
+    default_blocked = True
+
+    def __init__(self, vars=None, scaling=None, step_scale=.25, path_length=2., is_cov=False, step_rand=unif, state=None, model=None, **kwargs):
         """
         Parameters
         ----------
@@ -69,7 +71,7 @@ def __init__(self, vars=None, scaling=None, step_scale=.25, path_length=2., is_c
             state = SamplerHist()
         self.state = state
 
-        super(HamiltonianMC, self).__init__(vars, [model.fastlogp, model.fastdlogp(vars)])
+        super(HamiltonianMC, self).__init__(vars, [model.fastlogp, model.fastdlogp(vars)], **kwargs)
 
     def astep(self, q0, logp, dlogp):
         H = Hamiltonian(logp, dlogp, self.potential)

pymc/step_methods/metropolis.py

@@ -65,8 +65,10 @@ class Metropolis(ArrayStep):
         Optional model for sampling step. Defaults to None (taken from context).
 
     """
+    default_blocked = False
+
     def __init__(self, vars=None, S=None, proposal_dist=NormalProposal, scaling=1.,
-                 tune=True, tune_interval=100, model=None):
+                 tune=True, tune_interval=100, model=None, **kwargs):
 
         model = modelcontext(model)
 
@@ -84,7 +86,7 @@ def __init__(self, vars=None, S=None, proposal_dist=NormalProposal, scaling=1.,
         # Determine type of variables
         self.discrete = np.array([v.dtype in discrete_types for v in vars])
 
-        super(Metropolis, self).__init__(vars, [model.fastlogp])
+        super(Metropolis, self).__init__(vars, [model.fastlogp], **kwargs)
 
     def astep(self, q0, logp):
 

pymc/step_methods/nuts.py

@@ -17,13 +17,14 @@ class NUTS(ArrayStep):
     Hoffman, Matthew D., & Gelman, Andrew. (2011).
     The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo.
     """
+    default_blocked = True
     def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state=None,
                  Emax=1000,
                  target_accept=0.65,
                  gamma=0.05,
                  k=0.75,
                  t0=10,
-                 model=None):
+                 model=None, **kwargs):
         """
         Parameters
         ----------
@@ -83,7 +84,7 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state
 
 
 
-        super(NUTS, self).__init__(vars, [model.fastlogp, model.fastdlogp(vars)])
+        super(NUTS, self).__init__(vars, [model.fastlogp, model.fastdlogp(vars)], **kwargs)
 
     def astep(self, q0, logp, dlogp):
         H = Hamiltonian(logp, dlogp, self.potential)

pymc/step_methods/slicer.py

@@ -9,9 +9,9 @@
 
 
 class Slice(ArrayStep):
-
     """Slice sampler"""
-    def __init__(self, vars=None, w=1, tune=True, model=None):
+    default_blocked = False
+    def __init__(self, vars=None, w=1, tune=True, model=None, **kwargs):
 
         model = modelcontext(model)
 
@@ -23,7 +23,7 @@ def __init__(self, vars=None, w=1, tune=True, model=None):
         self.w_tune = []
         self.model = model
 
-        super(Slice, self).__init__(vars, [model.fastlogp])
+        super(Slice, self).__init__(vars, [model.fastlogp], **kwargs)
 
     def astep(self, q0, logp):
 

pymc/tests/test_plots.py

@@ -1,7 +1,7 @@
 import matplotlib
 matplotlib.use('Agg', warn=False)
 
-import numpy as np 
+import numpy as np
 from .checks import close_to
 
 import pymc.plots
@@ -58,7 +58,7 @@ def test_multichain_plots():
 
     autocorrplot(ptrace, vars=['switchpoint'])
 
-def test_make_2d(): 
+def test_make_2d():
 
     a = np.arange(4)
     close_to(pymc.plots.make_2d(a), a[:,None], 0)

pymc/tests/test_stats.py

@@ -56,7 +56,7 @@ def test_summary_0d_variable_model():
     tau = 1.3
     with Model() as model:
         x = Normal('x', mu, tau, testval=.1)
-        step = Metropolis(model.vars, np.diag([1.]))
+        step = Metropolis(model.vars, np.diag([1.]), blocked=True)
         trace = pm.sample(100, step=step)
     pm.summary(trace)
 
@@ -66,7 +66,7 @@ def test_summary_1d_variable_model():
     tau = 1.3
     with Model() as model:
         x = Normal('x', mu, tau, shape=2, testval=[.1, .1])
-        step = Metropolis(model.vars, np.diag([1.]))
+        step = Metropolis(model.vars, np.diag([1.]), blocked=True)
         trace = pm.sample(100, step=step)
     pm.summary(trace)
 
@@ -77,7 +77,7 @@ def test_summary_2d_variable_model():
     with Model() as model:
         x = Normal('x', mu, tau, shape=(2, 2),
                    testval=np.tile(.1, (2, 2)))
-        step = Metropolis(model.vars, np.diag([1.]))
+        step = Metropolis(model.vars, np.diag([1.]), blocked=True)
         trace = pm.sample(100, step=step)
     pm.summary(trace)
 

pymc/tests/test_step.py

@@ -1,5 +1,5 @@
 from .checks import *
-from .models import simple_model, mv_simple, mv_simple_discrete
+from .models import simple_model, mv_simple, mv_simple_discrete, simple_2model
 from theano.tensor import constant
 from scipy.stats.mstats import moment
 
@@ -13,15 +13,23 @@ def test_step_continuous():
     start, model, (mu, C) = mv_simple()
 
     with model:
-        hmc = pm.HamiltonianMC(scaling=C, is_cov=True)
-        mh = pm.Metropolis(S=C,
-                           proposal_dist=pm.MultivariateNormalProposal)
+        mh = pm.Metropolis()
         slicer = pm.Slice()
-        nuts = pm.NUTS(scaling=C, is_cov=True)
-        compound = pm.CompoundStep([hmc, mh])
+        hmc = pm.HamiltonianMC(scaling=C, is_cov=True, blocked=False)
+        nuts = pm.NUTS(scaling=C, is_cov=True, blocked=False)
+
+        mh_blocked = pm.Metropolis(S=C,
+                                   proposal_dist=pm.MultivariateNormalProposal,
+                                   blocked=True)
+        slicer_blocked = pm.Slice(blocked=True)
+        hmc_blocked = pm.HamiltonianMC(scaling=C, is_cov=True)
+        nuts_blocked = pm.NUTS(scaling=C, is_cov=True)
 
+        compound = pm.CompoundStep([hmc_blocked, mh_blocked])
 
-    steps = [mh, hmc, slicer, nuts, compound]
+
+    steps = [slicer, hmc, nuts, mh_blocked, hmc_blocked,
+             slicer_blocked, nuts_blocked, compound]
 
     unc = np.diag(C) ** .5
     check = [('x', np.mean, mu, unc / 10.),
@@ -32,6 +40,36 @@ def test_step_continuous():
         for (var, stat, val, bound) in check:
             yield check_stat, repr(st), h, var, stat, val, bound
 
+
+def test_non_blocked():
+    """Test that samplers correctly create non-blocked compound steps.
+    """
+
+    start, model = simple_2model()
+
+    with model:
+        # Metropolis and Slice are non-blocked by default
+        mh = pm.Metropolis()
+        assert isinstance(mh, pm.CompoundStep)
+        slicer = pm.Slice()
+        assert isinstance(slicer, pm.CompoundStep)
+        hmc = pm.HamiltonianMC(blocked=False)
+        assert isinstance(hmc, pm.CompoundStep)
+        nuts = pm.NUTS(blocked=False)
+        assert isinstance(nuts, pm.CompoundStep)
+
+        mh_blocked = pm.Metropolis(blocked=True)
+        assert isinstance(mh_blocked, pm.Metropolis)
+        slicer_blocked = pm.Slice(blocked=True)
+        assert isinstance(slicer_blocked, pm.Slice)
+        hmc_blocked = pm.HamiltonianMC()
+        assert isinstance(hmc_blocked, pm.HamiltonianMC)
+        nuts_blocked = pm.NUTS()
+        assert isinstance(nuts_blocked, pm.NUTS)
+
+        compound = pm.CompoundStep([hmc_blocked, mh_blocked])
+
+
 def test_step_discrete():
     start, model, (mu, C) = mv_simple_discrete()