Repository: pyro-ppl/pyro

Issue:
add Discrete distribution
this is useful for constructing marginals in particular.

it should be a very straightforward extension of Categorical (especially if implemented as a subclass?).


Discussion:
Should we swap the current implementations of Categorical with Discrete? Currently Categorical only takes `ps`, and `sample` returns index or a one-hot, while Discrete would take class labels (vector of indices) and would return a class label on `sample`, which is the opposite of say [webppl's implementation](https://github.com/probmods/webppl/blob/dev/src/dists.ad.js#L906)
yeah, that's probably why i thought i wanted Categorical originally. 

maybe we should combine them, having an optional `vs`, so that it can return index, 1-hot, or value depending on options?

Resolving patch:
diff --git a/pyro/distributions/categorical.py b/pyro/distributions/categorical.py
--- a/pyro/distributions/categorical.py
+++ b/pyro/distributions/categorical.py
@@ -2,6 +2,7 @@
 from torch.autograd import Variable
 from pyro.distributions.distribution import Distribution
 import itertools
+import numpy as np
 
 
 def _to_one_hot(x, ps):
@@ -24,38 +25,75 @@ class Categorical(Distribution):
     Categorical is a specialized version of multinomial where n = 1
     """
 
-    def __init__(self, ps, one_hot=True, batch_size=1, *args, **kwargs):
+    def __init__(self, ps, vs=None, one_hot=True, batch_size=1, *args, **kwargs):
         """
-        Constructor.
+        Instantiates a discrete distribution.
+        vs - tuple, list, numpy array, Variable, or torch tensor of values
+        ps - torch tensor of probabilities (must be same size as `vs`)
+        one_hot - return one-hot samples (when `vs` is None)
+        batch_size - expand ps and vs by a batch dimension
         """
+        if vs is not None:
+            if isinstance(vs, tuple):
+                # recursively turn tuples into lists
+                vs = [list(x) for x in vs]
+            if isinstance(vs, list):
+                vs = np.array(vs)
+            elif not isinstance(vs, Variable) and not isinstance(vs, torch.Tensor):
+                raise TypeError(("vs should be of type: list, Variable, Tensor, tuple, or numpy array"
+                                 "but was of {}".format(str(type(vs)))))
+        self.ps = ps
+        # vs is None, Variable(Tensor), or numpy.array
+        self.vs = vs
         if ps.dim() == 1:
             self.ps = ps.unsqueeze(0)
+            if isinstance(vs, Variable):
+                self.vs = vs.unsqueeze(0)
         elif batch_size > 1:
             self.ps = ps.unsqueeze(0).expand(batch_size, 0)
-        else:
-            self.ps = ps
-        # self.batch_size = batch_size
+            if isinstance(vs, Variable):
+                self.vs = vs.unsqueeze(0).expand(batch_size, 0)
         self.one_hot = one_hot
         super(Categorical, self).__init__(batch_size=1, *args, **kwargs)
 
     def sample(self):
         sample = Variable(torch.multinomial(self.ps.data, 1, replacement=True))
+        if self.vs is not None:
+            if isinstance(self.vs, np.ndarray):
+                # always returns a 2-d (unsqueezed 1-d) list
+                r = np.arange(self.vs.shape[0])
+                return [[x] for x in self.vs[r, sample.squeeze().data.numpy()].tolist()]
+            # self.vs is a torch.Tensor
+            return torch.gather(self.vs, 1, sample.long())
         if self.one_hot:
             # convert to onehot vector
             return _to_one_hot(sample, self.ps)
         return sample
 
     def batch_log_pdf(self, x, batch_size=1):
-        if x.dim() == 1 and batch_size == 1:
-            x = x.unsqueeze(0)
-        if self.one_hot:
+        if not isinstance(x, list):
+            if x.dim() == 1 and batch_size == 1:
+                x = x.unsqueeze(0)
             if self.ps.size(0) != x.size(0):
-                # convert to int torcho one-hot
+                # convert to int to one-hot
                 _ps = self.ps.expand_as(x)
             else:
                 _ps = self.ps
+        if self.vs is not None:
+            # get index of tensor
+            if isinstance(self.vs, np.ndarray):
+                # x is a list if self.vs was a list or np.array
+                bm = torch.Tensor((self.vs == np.array(x)).tolist())
+                ix = len(self.vs.shape) - 1
+                _x = torch.nonzero(bm).select(ix, ix)
+            else:
+                # x is a Variable(Tensor) as are ps and vs
+                ix = self.vs.dim() - 1
+                _x = torch.nonzero(self.vs.eq(x.expand_as(self.vs)).data).select(ix, ix)
+            x = Variable(_x).unsqueeze(1)
+        elif self.one_hot:
             return torch.sum(x * torch.log(_ps), 1)
-        return torch.log(torch.index_select(self.ps, 1, x.squeeze(1).long()))
+        return torch.log(torch.gather(self.ps, 1, x.long()))
 
     def log_pdf(self, _x):
         return torch.sum(self.batch_log_pdf(_x))
@@ -63,10 +101,29 @@ def log_pdf(self, _x):
     def support(self):
         r = self.ps.size(0)
         c = self.ps.size(1)
+
+        if self.vs is not None:
+            if isinstance(self.vs, np.ndarray):
+                # vs is an array, so the support must be of type array
+                r_np = self.vs.shape[0]
+                c_np = self.vs.shape[1]
+                ix = np.expand_dims(np.arange(r_np), axis=1)
+                b = torch.ones(r_np, 1)
+                return (self.vs[np.arange(r_np), torch.Tensor(list(x)).numpy().astype(int)]
+                        .reshape(r_np, 1).tolist()
+                        for x in itertools.product(torch.arange(0, c_np), repeat=r_np))
+            # vs is a tensor so support is of type tensor
+            return (torch.sum(self.vs * Variable(torch.Tensor(list(x))), 1)
+                    for x in itertools.product(torch.eye(c).numpy().tolist(),
+                    repeat=r))
+
         if self.one_hot:
-            return (Variable(torch.Tensor(list(x))) for x in itertools.product(torch.eye(c).numpy().tolist(),
+            return (Variable(torch.Tensor(list(x)))
+                    for x in itertools.product(torch.eye(c).numpy().tolist(),
                     repeat=r))
-        if self.ps.dim() == 1:
-            return iter([Variable(torch.Tensor([i])) for i in range(r)])
-        return (Variable(torch.Tensor(list(x)).unsqueeze(1)) for x in itertools.product(torch.arange(0, c),
+
+        if r == 1:
+            return iter([Variable(torch.Tensor([[i]])) for i in range(c)])
+        return (Variable(torch.Tensor(list(x)).unsqueeze(1))
+                for x in itertools.product(torch.arange(0, c),
                 repeat=r))


Test patch:
diff --git a/tests/test_data/support_categorical.json b/tests/test_data/support_categorical.json
--- a/tests/test_data/support_categorical.json
+++ b/tests/test_data/support_categorical.json
@@ -29,5 +29,27 @@
         [[2], [0]],
         [[2], [1]],
         [[2], [2]]
+    ],
+    "discrete": [
+        [[0], [3]],
+        [[0], [4]],
+        [[0], [5]],
+        [[1], [3]],
+        [[1], [4]],
+        [[1], [5]],
+        [[2], [3]],
+        [[2], [4]],
+        [[2], [5]]
+    ],
+    "discrete_arr": [
+        [["a"], ["d"]],
+        [["a"], ["e"]],
+        [["a"], ["f"]],
+        [["b"], ["d"]],
+        [["b"], ["e"]],
+        [["b"], ["f"]],
+        [["c"], ["d"]],
+        [["c"], ["e"]],
+        [["c"], ["f"]]
     ]
 }
diff --git a/tests/test_distributions.py b/tests/test_distributions.py
--- a/tests/test_distributions.py
+++ b/tests/test_distributions.py
@@ -157,12 +157,26 @@ def setUp(self):
         self.analytic_mean = n * self.ps
         one = Variable(torch.ones(3))
         self.analytic_var = n * torch.mul(self.ps, one.sub(self.ps))
+
+        # Discrete Distribution
+        self.d_ps = Variable(torch.Tensor([[0.2, 0.3, 0.5], [0.1, 0.1, 0.8]]))
+        self.d_vs = Variable(torch.Tensor([[0, 1, 2], [3, 4, 5]]))
+        self.d_vs_arr = [['a', 'b', 'c'], ['d', 'e', 'f']]
+        self.d_vs_tup = (('a', 'b', 'c'), ('d', 'e', 'f'))
+        self.d_test_data = Variable(torch.Tensor([[0], [5]]))
+        self.d_v_test_data = [['a'], ['f']]
+        self.d_dist = dist.Categorical(self.d_ps, self.d_vs)
+        self.d_dist_arr = dist.Categorical(self.d_ps, self.d_vs_arr)
+        self.d_dist_tup = dist.Categorical(self.d_ps, self.d_vs_tup)
+
         self.n_samples = 50000
 
         with open('tests/test_data/support_categorical.json') as data_file:
             data = json.load(data_file)
         self.support = list(map(lambda x: torch.Tensor(x), data['one_hot']))
         self.nhot_support = list(map(lambda x: torch.Tensor(x), data['not_hot']))
+        self.discrete_support = list(map(lambda x: torch.Tensor(x), data['discrete']))
+        self.discrete_arr_support = data['discrete_arr']
 
     def test_nhot_log_pdf(self):
         log_px_torch = self.dist_nhot.batch_log_pdf(self.test_data_nhot).data[0][0]
@@ -188,6 +202,31 @@ def test_mean_and_var(self):
         self.assertEqual(exp_, self.analytic_mean.data.numpy()[0], prec=0.01)
         self.assertEqual(torch_var, self.analytic_var.data.numpy()[0], prec=0.01)
 
+    def test_discrete_log_pdf(self):
+        log_px_torch = self.d_dist.batch_log_pdf(self.d_test_data).data[0][0]
+        log_px_np = float(spr.multinomial.logpmf(np.array([1, 0, 0]), 1, self.d_ps[0].data.numpy()))
+        log_px_torch2 = self.d_dist.batch_log_pdf(self.d_test_data).data[1][0]
+        log_px_np2 = float(spr.multinomial.logpmf(np.array([0, 0, 1]), 1, self.d_ps[1].data.numpy()))
+        self.assertEqual(log_px_torch, log_px_np, prec=1e-4)
+        self.assertEqual(log_px_torch2, log_px_np2, prec=1e-4)
+
+    def test_discrete_arr_logpdf(self):
+        log_px_torch = self.d_dist_arr.batch_log_pdf(self.d_v_test_data).data[0][0]
+        log_px_np = float(spr.multinomial.logpmf(np.array([1, 0, 0]), 1, self.d_ps[0].data.numpy()))
+        log_px_torch2 = self.d_dist.batch_log_pdf(self.d_test_data).data[1][0]
+        log_px_np2 = float(spr.multinomial.logpmf(np.array([0, 0, 1]), 1, self.d_ps[1].data.numpy()))
+        self.assertEqual(log_px_torch, log_px_np, prec=1e-4)
+        self.assertEqual(log_px_torch2, log_px_np2, prec=1e-4)
+
+    def test_discrete_support(self):
+        s = list(self.d_dist.support())
+        v = [torch.equal(x.data, y) for x, y in zip(s, self.discrete_support)]
+        self.assertTrue(all(v))
+
+    def test_discrete_arr_support(self):
+        s = list(self.d_dist_arr.support())
+        self.assertTrue(s == self.discrete_arr_support)
+
     def test_support(self):
         s = list(self.batch_dist.support())
         v = [torch.equal(x.data, y) for x, y in zip(s, self.support)]