Implemented CategoricalGibbsMetropolis, optimized BinaryGibbsMetropolis. (#1439)

* Implemented CategoricalGibbsMetropolis, optimized BinaryGibbsMetropolis.

* Added unit test for CategoricalGibbsMetropolis.
Also Added extra comments requested by PyMC3 devs, and fixed a minor bug
in the proportional proposal utilized by CategoricalGibbsMetropolis.

* Fix a minor bug in the newly added CategoricalGibbsMetropolis test.

Author: ozankabak

pymc3/sampling.py

@@ -9,7 +9,8 @@
 from .backends.ndarray import NDArray
 from .model import modelcontext, Point
 from .step_methods import (NUTS, HamiltonianMC, Metropolis, BinaryMetropolis,
-                           BinaryGibbsMetropolis, Slice, ElemwiseCategorical, CompoundStep)
+                           BinaryGibbsMetropolis, CategoricalGibbsMetropolis,
+                           Slice, CompoundStep)
 from tqdm import tqdm
 
 import sys
@@ -20,7 +21,7 @@
 
 def assign_step_methods(model, step=None, methods=(NUTS, HamiltonianMC, Metropolis,
                                                    BinaryMetropolis, BinaryGibbsMetropolis,
-                                                   Slice, ElemwiseCategorical)):
+                                                   Slice, CategoricalGibbsMetropolis)):
     '''
     Assign model variables to appropriate step methods. Passing a specified
     model will auto-assign its constituent stochastic variables to step methods

pymc3/step_methods/__init__.py

@@ -5,6 +5,7 @@
 from .metropolis import Metropolis
 from .metropolis import BinaryMetropolis
 from .metropolis import BinaryGibbsMetropolis
+from .metropolis import CategoricalGibbsMetropolis
 from .metropolis import NormalProposal
 from .metropolis import CauchyProposal
 from .metropolis import LaplaceProposal

pymc3/step_methods/gibbs.py

@@ -7,6 +7,7 @@
 from ..distributions.discrete import Categorical
 from numpy import array, max, exp, cumsum, nested_iters, empty, searchsorted, ones, arange
 from numpy.random import uniform
+from warnings import warn
 
 from theano.gof.graph import inputs
 from theano.tensor import add
@@ -25,6 +26,8 @@ class ElemwiseCategorical(ArrayStep):
     # variables)
 
     def __init__(self, vars, values=None, model=None):
+        warn('ElemwiseCategorical is deprecated, switch to CategoricalGibbsMetropolis.',
+             DeprecationWarning, stacklevel = 2)
         model = modelcontext(model)
         self.var = vars[0]
         self.sh = ones(self.var.dshape, self.var.dtype)
@@ -45,10 +48,7 @@ def competence(var):
         distribution = getattr(
             var.distribution, 'parent_dist', var.distribution)
         if isinstance(var.distribution, Categorical):
-            if var.distribution.k > 2:
-                return Competence.IDEAL
-            else:
-                return Competence.COMPATIBLE
+            return Competence.COMPATIBLE
         return Competence.INCOMPATIBLE
 
 

pymc3/step_methods/metropolis.py

@@ -2,12 +2,14 @@
 import numpy.random as nr
 import theano
 
+from ..distributions import draw_values
 from .arraystep import ArrayStepShared, ArrayStep, metrop_select, Competence
 import pymc3 as pm
 
 
-__all__ = ['Metropolis', 'BinaryMetropolis', 'BinaryGibbsMetropolis', 'NormalProposal',
-           'CauchyProposal', 'LaplaceProposal', 'PoissonProposal', 'MultivariateNormalProposal']
+__all__ = ['Metropolis', 'BinaryMetropolis', 'BinaryGibbsMetropolis',
+           'CategoricalGibbsMetropolis', 'NormalProposal', 'CauchyProposal',
+           'LaplaceProposal', 'PoissonProposal', 'MultivariateNormalProposal']
 
 # Available proposal distributions for Metropolis
 
@@ -239,41 +241,51 @@ def competence(var):
 
 
 class BinaryGibbsMetropolis(ArrayStep):
-    """Metropolis-Hastings optimized for binary variables"""
+    """A Metropolis-within-Gibbs step method optimized for binary variables"""
 
     def __init__(self, vars, order='random', model=None):
 
         model = pm.modelcontext(model)
 
         self.dim = sum(v.dsize for v in vars)
-        self.order = order
+
+        if order == 'random':
+            self.shuffle_dims = True
+            self.order = list(range(self.dim))
+        else:
+            if sorted(order) != list(range(self.dim)):
+                raise ValueError('Argument \'order\' has to be a permutation')
+            self.shuffle_dims = False
+            self.order = order
 
         if not all([v.dtype in pm.discrete_types for v in vars]):
             raise ValueError(
-                'All variables must be Bernoulli for BinaryGibbsMetropolis')
+                'All variables must be binary for BinaryGibbsMetropolis')
 
         super(BinaryGibbsMetropolis, self).__init__(vars, [model.fastlogp])
 
     def astep(self, q0, logp):
-        order = list(range(self.dim))
-        if self.order == 'random':
+        order = self.order
+        if self.shuffle_dims:
             nr.shuffle(order)
 
-        q_prop = np.copy(q0)
-        q_cur = np.copy(q0)
+        q = np.copy(q0)
+        logp_curr = logp(q)
 
         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 = np.copy(q_cur)
+            curr_val, q[idx] = q[idx], True - q[idx]
+            logp_prop = logp(q)
+            q[idx] = metrop_select(logp_prop - logp_curr, q[idx], curr_val)
+            if q[idx] != curr_val:
+                logp_curr = logp_prop
 
-        return q_cur
+        return q
 
     @staticmethod
     def competence(var):
         '''
-        BinaryMetropolis is only suitable for binary (bool)
-        and Categorical variables with k=1.
+        BinaryMetropolis is only suitable for Bernoulli
+        and Categorical variables with k=2.
         '''
         distribution = getattr(
             var.distribution, 'parent_dist', var.distribution)
@@ -283,6 +295,132 @@ def competence(var):
             return Competence.IDEAL
         return Competence.INCOMPATIBLE
 
+class CategoricalGibbsMetropolis(ArrayStep):
+    """A Metropolis-within-Gibbs step method optimized for categorical variables.
+       This step method works for Bernoulli variables as well, but it is not
+       optimized for them, like BinaryGibbsMetropolis is. Step method supports
+       two types of proposals: A uniform proposal and a proportional proposal,
+       which was introduced by Liu in his 1996 technical report
+       "Metropolized Gibbs Sampler: An Improvement".
+    """
+
+    def __init__(self, vars, proposal='uniform', order='random', model=None):
+
+        model = pm.modelcontext(model)
+        vars = pm.inputvars(vars)
+
+        dimcats = []
+        # The above variable is a list of pairs (aggregate dimension, number
+        # of categories). For example, if vars = [x, y] with x being a 2-D
+        # variable with M categories and y being a 3-D variable with N
+        # categories, we will have dimcats = [(0, M), (1, M), (2, N), (3, N), (4, N)].
+        for v in vars:
+            distr = getattr(v.distribution, 'parent_dist', v.distribution)
+            if isinstance(distr, pm.Categorical):
+                k = draw_values([distr.k])[0]
+            elif isinstance(distr, pm.Bernoulli) or (v.dtype in pm.bool_types):
+                k = 2
+            else:
+                raise ValueError('All variables must be categorical or binary' +
+                                 'for CategoricalGibbsMetropolis')
+            start = len(dimcats)
+            dimcats += [(dim, k) for dim in range(start, start + v.dsize)]
+
+        if order == 'random':
+            self.shuffle_dims = True
+            self.dimcats = dimcats
+        else:
+            if sorted(order) != list(range(len(dimcats))):
+                raise ValueError('Argument \'order\' has to be a permutation')
+            self.shuffle_dims = False
+            self.dimcats = [dimcats[j] for j in order]
+
+        if proposal == 'uniform':
+            self.astep = self.astep_unif
+        elif proposal == 'proportional':
+            # Use the optimized "Metropolized Gibbs Sampler" described in Liu96.
+            self.astep = self.astep_prop
+        else:
+            raise ValueError('Argument \'proposal\' should either be ' +
+                    '\'uniform\' or \'proportional\'')
+
+        super(CategoricalGibbsMetropolis, self).__init__(vars, [model.fastlogp])
+
+    def astep_unif(self, q0, logp):
+        dimcats = self.dimcats
+        if self.shuffle_dims:
+            nr.shuffle(dimcats)
+
+        q = np.copy(q0)
+        logp_curr = logp(q)
+
+        for dim, k in dimcats:
+            curr_val, q[dim] = q[dim], sample_except(k, q[dim])
+            logp_prop = logp(q)
+            q[dim] = metrop_select(logp_prop - logp_curr, q[dim], curr_val)
+            if q[dim] != curr_val:
+                logp_curr = logp_prop
+
+        return q
+
+    def astep_prop(self, q0, logp):
+        dimcats = self.dimcats
+        if self.shuffle_dims:
+            nr.shuffle(dimcats)
+
+        q = np.copy(q0)
+        logp_curr = logp(q)
+
+        for dim, k in dimcats:
+            logp_curr = self.metropolis_proportional(q, logp, logp_curr, dim, k)
+
+        return q
+
+    def metropolis_proportional(self, q, logp, logp_curr, dim, k):
+        given_cat = int(q[dim])
+        log_probs = np.zeros(k)
+        log_probs[given_cat] = logp_curr
+        candidates = list(range(k))
+        for candidate_cat in candidates:
+            if candidate_cat != given_cat:
+                q[dim] = candidate_cat
+                log_probs[candidate_cat] = logp(q)
+        probs = softmax(log_probs)
+        prob_curr, probs[given_cat] = probs[given_cat], 0.0
+        probs /= (1.0 - prob_curr)
+        proposed_cat = nr.choice(candidates, p = probs)
+        accept_ratio = (1.0 - prob_curr) / (1.0 - probs[proposed_cat])
+        if not np.isfinite(accept_ratio) or nr.uniform() >= accept_ratio:
+            q[dim] = given_cat
+            return logp_curr
+        q[dim] = proposed_cat
+        return log_probs[proposed_cat]
+
+    @staticmethod
+    def competence(var):
+        '''
+        CategoricalGibbsMetropolis is only suitable for Bernoulli and
+        Categorical variables.
+        '''
+        distribution = getattr(
+            var.distribution, 'parent_dist', var.distribution)
+        if isinstance(distribution, pm.Categorical):
+            if distribution.k > 2:
+                return Competence.IDEAL
+            return Competence.COMPATIBLE
+        elif isinstance(distribution, pm.Bernoulli) or (var.dtype in pm.bool_types):
+            return Competence.COMPATIBLE
+        return Competence.INCOMPATIBLE
+
+def sample_except(limit, excluded):
+    candidate = nr.choice(limit - 1)
+    if candidate >= excluded:
+        candidate += 1
+    return candidate
+
+def softmax(x):
+    e_x = np.exp(x - np.max(x))
+    return e_x / np.sum(e_x, axis = 0)
 
 def delta_logp(logp, vars, shared):
     [logp0], inarray0 = pm.join_nonshared_inputs([logp], vars, shared)

pymc3/tests/models.py

@@ -1,4 +1,4 @@
-from pymc3 import Model, Normal, Metropolis
+from pymc3 import Model, Normal, Categorical, Metropolis
 import numpy as np
 import pymc3 as pm
 from itertools import product
@@ -14,6 +14,17 @@ def simple_model():
     return model.test_point, model, (mu, tau ** -1)
 
 
+def simple_categorical():
+    p = np.array([0.1, 0.2, 0.3, 0.4])
+    v = np.array([0.0, 1.0, 2.0, 3.0])
+    with Model() as model:
+        Categorical('x', p, shape = 3, testval = [1, 2, 3])
+
+    mu = np.dot(p, v)
+    var = np.dot(p, (v - mu) ** 2)
+    return model.test_point, model, (mu, var)
+
+
 def multidimensional_model():
     mu = -2.1
     tau = 1.3

pymc3/tests/test_step.py

@@ -1,11 +1,12 @@
 import unittest
 
 from .checks import close_to
-from .models import mv_simple, mv_simple_discrete, simple_2model
+from .models import simple_categorical, 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, BinaryGibbsMetropolis, Metropolis, Constant, Slice,
-                                CompoundStep, MultivariateNormalProposal, HamiltonianMC)
+from pymc3.step_methods import (NUTS, BinaryGibbsMetropolis, CategoricalGibbsMetropolis,
+                                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
@@ -52,6 +53,21 @@ def test_step_discrete(self):
             trace = sample(20000, step=step, start=start, model=model, random_seed=1)
             self.check_stat(check, trace)
 
+    def test_step_categorical(self):
+        start, model, (mu, C) = simple_categorical()
+        unc = C ** .5
+        check = (('x', np.mean, mu, unc / 10.),
+                 ('x', np.std, unc, unc / 10.))
+        with model:
+            steps = (
+                CategoricalGibbsMetropolis(model.x, proposal = 'uniform'),
+                CategoricalGibbsMetropolis(model.x, proposal = 'proportional'),
+            )
+        for step in steps:
+            trace = sample(8000, 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)
@@ -95,17 +111,15 @@ def test_normal(self):
         self.assertIsInstance(steps, NUTS)
 
     def test_categorical(self):
-        """Test categorical distribution is assigned binary gibbs metropolis method"""
+        """Test categorical distribution is assigned categorical 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:
+            Categorical('y', np.array([0.25, 0.70, 0.05]))
+            steps = assign_step_methods(model, [])
+        self.assertIsInstance(steps, CategoricalGibbsMetropolis)
 
     def test_binomial(self):
         """Test binomial distribution is assigned metropolis method."""