Generalized multinomial to accept multiple observations (#1390)

* Generalized multinomial to accept multiple observations

* Added test for multinomial with multiple observations

* Added shape to Multinomial call in test_multinomial_vec

* Fix for test failur in multinomial

* Fix for test failure in multinomial

* Added correct shape to multinomial_vec test

* Dialed down assert_equal precision in mutlinomial test

Author: Chris Fonnesbeck

pymc3/distributions/multivariate.py

@@ -257,10 +257,20 @@ def random(self, point=None, size=None):
     def logp(self, x):
         n = self.n
         p = self.p
-        # only defined for sum(p) == 1
+        
+        if x.ndim==2:
+            x_sum = x.sum(axis=0)
+            n_sum = n * x.shape[0]
+        else:
+            x_sum = x
+            n_sum = n
+
         return bound(
-            factln(n) + tt.sum(x * tt.log(p) - factln(x)),
-            tt.all(x >= 0), tt.all(x <= n), tt.eq(tt.sum(x), n),
+            factln(n_sum) + tt.sum(x_sum * tt.log(p) - factln(x_sum)),
+            tt.all(x >= 0), 
+            tt.all(x <= n), 
+            tt.eq(tt.sum(x_sum), n_sum),
+            tt.isclose(p.sum(), 1),
             n >= 0)
 
 

pymc3/tests/test_distributions.py

@@ -560,6 +560,18 @@ def check_multinomial(self, n):
         self.pymc3_matches_scipy(Multinomial, Vector(Nat, n), {'p': Simplex(n), 'n': Nat},
                                  multinomial_logpdf)
 
+    def test_multinomial_vec(self):
+        vals = np.array([[2,4,4], [3,3,4]])
+        p = np.array([0.2, 0.3, 0.5])
+        n = 10
+        with Model() as model:
+            Multinomial('m', n=10, p=p, shape=vals.shape)
+        pt = {'m': vals}
+        with Model() as model_sum:
+            Multinomial('m_sum', n=2*n, p=p, shape=len(p))
+        pt_sum = {'m_sum': vals.sum(0)}
+        assert_almost_equal(model.fastlogp(pt), model_sum.fastlogp(pt_sum), decimal=4)
+
     def test_categorical(self):
         for n in [2, 3, 4]:
             yield self.check_categorical, n