Refactor assign step methods (pymc-devs#1347)

* Refactor step tests

* Refactor assign_step_methods and test_step

Author: ColCarroll

pymc3/sampling.py

@@ -2,13 +2,12 @@
 from .backends.base import merge_traces, BaseTrace, MultiTrace
 from .backends.ndarray import NDArray
 from joblib import Parallel, delayed
-from time import time
 from .model import modelcontext, Point
 from .step_methods import (NUTS, HamiltonianMC, Metropolis, BinaryMetropolis,
                            BinaryGibbsMetropolis, Slice, ElemwiseCategorical, CompoundStep)
 from .progressbar import progress_bar
 from numpy.random import randint, seed
-from numpy import shape, append, asarray
+from numpy import shape, asarray
 from collections import defaultdict
 
 import sys
@@ -17,9 +16,9 @@
 __all__ = ['sample', 'iter_sample', 'sample_ppc']
 
 
-def assign_step_methods(model, step=None,
-                        methods=(NUTS, HamiltonianMC, Metropolis, BinaryMetropolis, BinaryGibbsMetropolis,
-                                 Slice, ElemwiseCategorical)):
+def assign_step_methods(model, step=None, methods=(NUTS, HamiltonianMC, Metropolis,
+                                                   BinaryMetropolis, BinaryGibbsMetropolis,
+                                                   Slice, ElemwiseCategorical)):
     '''
     Assign model variables to appropriate step methods. Passing a specified
     model will auto-assign its constituent stochastic variables to step methods
@@ -50,31 +49,29 @@ def assign_step_methods(model, step=None,
     steps = []
     assigned_vars = set()
     if step is not None:
-        steps = append(steps, step).tolist()
-        for s in steps:
+        try:
+            steps += list(step)
+        except TypeError:
+            steps.append(step)
+        for step in steps:
             try:
-                assigned_vars = assigned_vars | set(s.vars)
+                assigned_vars = assigned_vars.union(set(step.vars))
             except AttributeError:
-                for m in s.methods:
-                    assigned_vars = assigned_vars | set(m.vars)
+                for method in step.methods:
+                    assigned_vars = assigned_vars.union(set(method.vars))
 
     # Use competence classmethods to select step methods for remaining
     # variables
     selected_steps = defaultdict(list)
     for var in model.free_RVs:
-        if not var in assigned_vars:
-
-            competences = {s: s._competence(var) for s in methods}
-
-            selected = max(competences.keys(), key=(lambda k: competences[k]))
-
+        if var not in assigned_vars:
+            selected = max(methods, key=lambda method: method._competence(var))
             if model.verbose:
                 print('Assigned {0} to {1}'.format(selected.__name__, var))
             selected_steps[selected].append(var)
 
     # Instantiate all selected step methods
-    steps += [s(vars=selected_steps[s])
-              for s in selected_steps if selected_steps[s]]
+    steps += [step(vars=selected_steps[step]) for step in selected_steps if selected_steps[step]]
 
     if len(steps) == 1:
         steps = steps[0]

pymc3/step_methods/arraystep.py

@@ -105,9 +105,9 @@ def __init__(self, vars, fs, allvars=False, blocked=True):
     def step(self, point):
         bij = DictToArrayBijection(self.ordering, point)
 
-        inputs = list(map(bij.mapf, self.fs))
+        inputs = [bij.mapf(x) for x in self.fs]
         if self.allvars:
-            inputs += [point]
+            inputs.append(point)
 
         apoint = self.astep(bij.map(point), *inputs)
         return bij.rmap(apoint)

pymc3/tests/test_step.py

@@ -1,139 +1,115 @@
+import unittest
+
 from .checks import close_to
-from .models import simple_model, mv_simple, mv_simple_discrete, simple_2model
-from ..step_methods import MultivariateNormalProposal
-from theano.tensor import constant
-from scipy.stats.mstats import moment
+from .models import mv_simple, mv_simple_discrete, simple_2model
 from pymc3.sampling import assign_step_methods, sample
 from pymc3.model import Model
-from pymc3.step_methods import NUTS, BinaryMetropolis, BinaryGibbsMetropolis, Metropolis, Constant, ElemwiseCategorical, Slice, CompoundStep, MultivariateNormalProposal, HamiltonianMC
+from pymc3.step_methods import (NUTS, BinaryGibbsMetropolis, Metropolis, Constant, Slice,
+                                CompoundStep, MultivariateNormalProposal, HamiltonianMC)
 from pymc3.distributions import Binomial, Normal, Bernoulli, Categorical
 from numpy.testing import assert_almost_equal
 import numpy as np
 
 
-def check_stat(name, trace, var, stat, value, bound):
-    s = stat(trace[var][2000:], axis=0)
-    close_to(s, value, bound)
-
-
-def test_step_continuous():
-    start, model, (mu, C) = mv_simple()
-
-    with model:
-        mh = Metropolis()
-        slicer = Slice()
-        hmc = HamiltonianMC(scaling=C, is_cov=True, blocked=False)
-        nuts = NUTS(scaling=C, is_cov=True, blocked=False)
-
-        mh_blocked = Metropolis(S=C,
-                                proposal_dist=MultivariateNormalProposal,
-                                blocked=True)
-        slicer_blocked = Slice(blocked=True)
-        hmc_blocked = HamiltonianMC(scaling=C, is_cov=True)
-        nuts_blocked = NUTS(scaling=C, is_cov=True)
-
-        compound = CompoundStep([hmc_blocked, mh_blocked])
-
-    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.),
-             ('x', np.std, unc, unc / 10.)]
-
-    for st in steps:
-        h = sample(8000, st, start, model=model, random_seed=1)
-        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 = Metropolis()
-        assert isinstance(mh, CompoundStep)
-        slicer = Slice()
-        assert isinstance(slicer, CompoundStep)
-        hmc = HamiltonianMC(blocked=False)
-        assert isinstance(hmc, CompoundStep)
-        nuts = NUTS(blocked=False)
-        assert isinstance(nuts, CompoundStep)
-
-        mh_blocked = Metropolis(blocked=True)
-        assert isinstance(mh_blocked, Metropolis)
-        slicer_blocked = Slice(blocked=True)
-        assert isinstance(slicer_blocked, Slice)
-        hmc_blocked = HamiltonianMC()
-        assert isinstance(hmc_blocked, HamiltonianMC)
-        nuts_blocked = NUTS()
-        assert isinstance(nuts_blocked, NUTS)
-
-        compound = CompoundStep([hmc_blocked, mh_blocked])
-
-
-def test_step_discrete():
-    start, model, (mu, C) = mv_simple_discrete()
-
-    with model:
-        mh = Metropolis(S=C,
-                        proposal_dist=MultivariateNormalProposal)
-        slicer = Slice()
-
-    steps = [mh]
-
-    unc = np.diag(C) ** .5
-    check = [('x', np.mean, mu, unc / 10.),
-             ('x', np.std, unc, unc / 10.)]
-
-    for st in steps:
-        h = sample(20000, st, start, model=model, random_seed=1)
-
-        for (var, stat, val, bound) in check:
-            yield check_stat, repr(st), h, var, stat, val, bound
-
-
-def test_constant_step():
-
-    with Model() as model:
-        x = Normal('x', 0, 1)
-        start = {'x': -1}
-        tr = sample(10, step=Constant([x]), start=start)
+class TestStepMethods(object):  # yield test doesn't work subclassing unittest.TestCase
+    def check_stat(self, check, trace):
+        for (var, stat, value, bound) in check:
+            s = stat(trace[var][2000:], axis=0)
+            close_to(s, value, bound)
+
+    def test_step_continuous(self):
+        start, model, (mu, C) = mv_simple()
+        unc = np.diag(C) ** .5
+        check = (('x', np.mean, mu, unc / 10.),
+                 ('x', np.std, unc, unc / 10.))
+        with model:
+            steps = (
+                Slice(),
+                HamiltonianMC(scaling=C, is_cov=True, blocked=False),
+                NUTS(scaling=C, is_cov=True, blocked=False),
+                Metropolis(S=C, proposal_dist=MultivariateNormalProposal, blocked=True),
+                Slice(blocked=True),
+                HamiltonianMC(scaling=C, is_cov=True),
+                NUTS(scaling=C, is_cov=True),
+                CompoundStep([
+                    HamiltonianMC(scaling=C, is_cov=True),
+                    HamiltonianMC(scaling=C, is_cov=True, blocked=False)]),
+            )
+        for step in steps:
+            trace = sample(8000, step=step, start=start, model=model, random_seed=1)
+            yield self.check_stat, check, trace
+
+    def test_step_discrete(self):
+        start, model, (mu, C) = mv_simple_discrete()
+        unc = np.diag(C) ** .5
+        check = (('x', np.mean, mu, unc / 10.),
+                 ('x', np.std, unc, unc / 10.))
+        with model:
+            steps = (
+                Metropolis(S=C, proposal_dist=MultivariateNormalProposal),
+            )
+        for step in steps:
+            trace = sample(20000, step=step, start=start, model=model, random_seed=1)
+            self.check_stat(check, trace)
+
+    def test_constant_step(self):
+        with Model():
+            x = Normal('x', 0, 1)
+            start = {'x': -1}
+            tr = sample(10, step=Constant([x]), start=start)
         assert_almost_equal(tr['x'], start['x'], decimal=10)
 
 
-def test_assign_step_methods():
-
-    with Model() as model:
-        x = Bernoulli('x', 0.5)
-        steps = assign_step_methods(model, [])
-
-        assert isinstance(steps, BinaryGibbsMetropolis)
-
-    with Model() as model:
-        x = Normal('x', 0, 1)
-        steps = assign_step_methods(model, [])
-
-        assert isinstance(steps, NUTS)
-
-    with Model() as model:
-        x = Categorical('x', np.array([0.25, 0.75]))
-        steps = assign_step_methods(model, [])
-
-        assert isinstance(steps, BinaryGibbsMetropolis)
+class TestCompoundStep(unittest.TestCase):
+    samplers = (Metropolis, Slice, HamiltonianMC, NUTS)
+
+    def test_non_blocked(self):
+        """Test that samplers correctly create non-blocked compound steps."""
+        _, model = simple_2model()
+        with model:
+            for sampler in self.samplers:
+                self.assertIsInstance(sampler(blocked=False), CompoundStep)
+
+    def test_blocked(self):
+        _, model = simple_2model()
+        with model:
+            for sampler in self.samplers:
+                sampler_instance = sampler(blocked=True)
+                self.assertNotIsInstance(sampler_instance, CompoundStep)
+                self.assertIsInstance(sampler_instance, sampler)
+
+
+class TestAssignStepMethods(unittest.TestCase):
+    def test_bernoulli(self):
+        """Test bernoulli distribution is assigned binary gibbs metropolis method"""
+        with Model() as model:
+            Bernoulli('x', 0.5)
+            steps = assign_step_methods(model, [])
+        self.assertIsInstance(steps, BinaryGibbsMetropolis)
+
+    def test_normal(self):
+        """Test normal distribution is assigned NUTS method"""
+        with Model() as model:
+            Normal('x', 0, 1)
+            steps = assign_step_methods(model, [])
+        self.assertIsInstance(steps, NUTS)
+
+    def test_categorical(self):
+        """Test categorical distribution is assigned binary gibbs metropolis method"""
+        with Model() as model:
+            Categorical('x', np.array([0.25, 0.75]))
+            steps = assign_step_methods(model, [])
+        self.assertIsInstance(steps, BinaryGibbsMetropolis)
 
     # with Model() as model:
     #     x = Categorical('x', np.array([0.25, 0.70, 0.05]))
     #     steps = assign_step_methods(model, [])
     #
     #     assert isinstance(steps, ElemwiseCategoricalStep)
 
-    with Model() as model:
-        x = Binomial('x', 10, 0.5)
-        steps = assign_step_methods(model, [])
-
-        assert isinstance(steps, Metropolis)
+    def test_binomial(self):
+        """Test binomial distribution is assigned metropolis method."""
+        with Model() as model:
+            Binomial('x', 10, 0.5)
+            steps = assign_step_methods(model, [])
+        self.assertIsInstance(steps, Metropolis)