Merge pull request #799 from pymc-devs/gibbs

Add Binary Gibbs Metropolis sampler for mixture models.

Author: springcoil

pymc3/sampling.py

@@ -12,21 +12,21 @@
 __all__ = ['sample', 'iter_sample', 'sample_ppc']
 
 def assign_step_methods(model, step=None,
-        methods=(NUTS, HamiltonianMC, Metropolis, BinaryMetropolis, 
+        methods=(NUTS, HamiltonianMC, Metropolis, BinaryMetropolis, BinaryGibbsMetropolis,
         Slice, ElemwiseCategoricalStep)):
     '''
-    Assign model variables to appropriate step methods. Passing a specified 
+    Assign model variables to appropriate step methods. Passing a specified
     model will auto-assign its constituent stochastic variables to step methods
     based on the characteristics of the variables. This function is intended to
     be called automatically from `sample()`, but may be called manually. Each
     step method passed should have a `competence()` method that returns an
     ordinal competence value corresponding to the variable passed to it. This
-    value quantifies the appropriateness of the step method for sampling the 
-    variable. 
-    
+    value quantifies the appropriateness of the step method for sampling the
+    variable.
+
     Parameters
     ----------
-    
+
     model : Model object
         A fully-specified model object
     step : step function or vector of step functions
@@ -35,12 +35,12 @@ def assign_step_methods(model, step=None,
     methods : vector of step method classes
         The set of step methods from which the function may choose. Defaults
         to the main step methods provided by PyMC3.
-        
+
     Returns
     -------
     List of step methods associated with the model's variables.
     '''
-        
+
     steps = []
     assigned_vars = set()
     if step is not None:
@@ -51,26 +51,26 @@ def assign_step_methods(model, step=None,
             except AttributeError:
                 for m in s.methods:
                     assigned_vars = assigned_vars | set(m.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 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]]
-    
+
     if len(steps)==1:
-        steps = steps[0]    
-                
+        steps = steps[0]
+
     return steps
 
 def sample(draws, step=None, start=None, trace=None, chain=0, njobs=1, tune=None,
@@ -86,8 +86,8 @@ def sample(draws, step=None, start=None, trace=None, chain=0, njobs=1, tune=None
     draws : int
         The number of samples to draw
     step : function or iterable of functions
-        A step function or collection of functions. If no step methods are 
-        specified, or are partially specified, they will be assigned 
+        A step function or collection of functions. If no step methods are
+        specified, or are partially specified, they will be assigned
         automatically (defaults to None).
     start : dict
         Starting point in parameter space (or partial point)
@@ -120,7 +120,7 @@ def sample(draws, step=None, start=None, trace=None, chain=0, njobs=1, tune=None
     MultiTrace object with access to sampling values
     """
     model = modelcontext(model)
-    
+
     step = assign_step_methods(model, step)
 
     if njobs is None:
@@ -184,7 +184,7 @@ def iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
 
     draws : int
         The number of samples to draw
-    step : function 
+    step : function
         Step function
     start : dict
         Starting point in parameter space (or partial point)

pymc3/step_methods/metropolis.py

@@ -13,7 +13,7 @@
 from ..theanof import make_shared_replacements, join_nonshared_inputs, CallableTensor
 
 
-__all__ = ['Metropolis', 'BinaryMetropolis', 'NormalProposal', 'CauchyProposal', 'LaplaceProposal', 'PoissonProposal', 'MultivariateNormalProposal']
+__all__ = ['Metropolis', 'BinaryMetropolis', 'BinaryGibbsMetropolis', 'NormalProposal', 'CauchyProposal', 'LaplaceProposal', 'PoissonProposal', 'MultivariateNormalProposal']
 
 # Available proposal distributions for Metropolis
 
@@ -71,7 +71,7 @@ class Metropolis(ArrayStepShared):
 
     """
     default_blocked = False
-    
+
     def __init__(self, vars=None, S=None, proposal_dist=NormalProposal, scaling=1.,
                  tune=True, tune_interval=100, model=None, **kwargs):
 
@@ -110,6 +110,7 @@ def astep(self, q0):
             self.accepted = 0
 
         delta = self.proposal_dist() * self.scaling
+
         if self.any_discrete:
             if self.all_discrete:
                 delta = round(delta, 0).astype(int)
@@ -121,16 +122,15 @@ def astep(self, q0):
         else:
             q = q0 + delta
 
-
-        q_new = metrop_select(self.delta_logp(q,q0), q, q0)
+        q_new = metrop_select(self.delta_logp(q, q0), q, q0)
 
         if q_new is q:
             self.accepted += 1
 
         self.steps_until_tune -= 1
 
         return q_new
-        
+
     @staticmethod
     def competence(var):
         if var.dtype in discrete_types:
@@ -179,7 +179,7 @@ def tune(scale, acc_rate):
 
 class BinaryMetropolis(ArrayStep):
     """Metropolis-Hastings optimized for binary variables"""
-    
+
     def __init__(self, vars, scaling=1., tune=True, tune_interval=100, model=None):
 
         model = modelcontext(model)
@@ -197,31 +197,71 @@ def __init__(self, vars, scaling=1., tune=True, tune_interval=100, model=None):
         super(BinaryMetropolis, self).__init__(vars, [model.fastlogp])
 
     def astep(self, q0, logp):
-
         # Convert adaptive_scale_factor to a jump probability
         p_jump = 1. - .5 ** self.scaling
 
         rand_array = random(q0.shape)
         q = copy(q0)
         # Locations where switches occur, according to p_jump
-        switch_locs = where(rand_array < p_jump)
+        switch_locs = (rand_array < p_jump)
         q[switch_locs] = True - q[switch_locs]
-
         q_new = metrop_select(logp(q) - logp(q0), q, q0)
 
         return q_new
-        
+
+    @staticmethod
+    def competence(var):
+        '''
+        BinaryMetropolis is only suitable for binary (bool)
+        and Categorical variables with k=1.
+        '''
+        if isinstance(var.distribution, Bernoulli) or (var.dtype in bool_types):
+            return Competence.compatible
+        elif isinstance(var.distribution, Categorical) and (var.distribution.k == 2):
+            return Competence.compatible
+        return Competence.incompatible
+
+class BinaryGibbsMetropolis(ArrayStep):
+    """Metropolis-Hastings optimized for binary variables"""
+
+    def __init__(self, vars, order='random', model=None):
+
+        model = modelcontext(model)
+
+        self.dim = sum(v.dsize for v in vars)
+        self.order = order
+
+        if not all([v.dtype in discrete_types for v in vars]):
+            raise ValueError(
+                'All variables must be Bernoulli for BinaryGibbsMetropolis')
+
+        super(BinaryGibbsMetropolis, self).__init__(vars, [model.fastlogp])
+
+    def astep(self, q0, logp):
+        order = list(range(self.dim))
+        if self.order == 'random':
+            np.random.shuffle(order)
+
+        q_prop = copy(q0)
+        q_cur = copy(q0)
+
+        for idx in order:
+            q_prop[idx] = True - q_prop[idx]
+            q_cur = metrop_select(logp(q_prop) - logp(q_cur), q_prop, q_cur)
+            q_prop = copy(q_cur)
+
+        return q_cur
+
     @staticmethod
     def competence(var):
         '''
-        BinaryMetropolis is only suitable for binary (bool) 
+        BinaryMetropolis is only suitable for binary (bool)
         and Categorical variables with k=1.
         '''
         if isinstance(var.distribution, Bernoulli) or (var.dtype in bool_types):
             return Competence.ideal
-        if isinstance(var.distribution, Categorical):
-            if var.distribution.k==2:
-                return Competence.ideal
+        elif isinstance(var.distribution, Categorical) and (var.distribution.k == 2):
+            return Competence.ideal
         return Competence.incompatible
 
 def delta_logp(logp, vars, shared):

pymc3/tests/test_step.py

@@ -4,7 +4,7 @@
 from scipy.stats.mstats import moment
 from pymc3.sampling import assign_step_methods
 from pymc3.model import Model
-from pymc3.step_methods import NUTS, BinaryMetropolis, Metropolis, Constant, ElemwiseCategoricalStep
+from pymc3.step_methods import NUTS, BinaryMetropolis, BinaryGibbsMetropolis, Metropolis, Constant, ElemwiseCategoricalStep
 from pymc3.distributions import Binomial, Normal, Bernoulli, Categorical
 from numpy.testing import assert_almost_equal
 
@@ -96,41 +96,41 @@ def test_step_discrete():
             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)
         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, BinaryMetropolis)
-    
+
+        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, BinaryMetropolis)
-        
+
+        assert isinstance(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)
+
+        assert isinstance(steps, Metropolis)