Symbolic dimensions can be passed to Multinomial, NegativeMultinomial

sympy/stats/joint_rv.py

@@ -101,7 +101,6 @@ def marginal_distribution(self, *indices):
         if self.distribution.is_Continuous:
             f = Lambda(sym, integrate(self.distribution(*all_syms), *limits))
         elif self.distribution.is_Discrete:
-            limits = [(limit[0], limit[1].inf, limit[1].sup) for limit in limits]
             f = Lambda(sym, summation(self.distribution(*all_syms), *limits))
         return f.xreplace(replace_dict)
 

sympy/stats/joint_rv_types.py

@@ -1,6 +1,6 @@
 from sympy import (sympify, S, pi, sqrt, exp, Lambda, Indexed, Gt, IndexedBase,
                     besselk, gamma, Interval, Range, factorial, Mul, Integer,
-                    Add, rf, Eq, Piecewise)
+                    Add, rf, Eq, Piecewise, Symbol, imageset)
 from sympy.matrices import ImmutableMatrix
 from sympy.matrices.expressions.determinant import det
 from sympy.stats.joint_rv import (JointDistribution, JointPSpace,
@@ -404,21 +404,24 @@ class MultinomialDistribution(JointDistribution):
     is_Discrete = True
 
     def check(self, n, p):
-        _value_check(((n > 0) != False) and isinstance(n, Integer),
+        _value_check(((n > 0) != False),
                         "number of trials must be a positve integer")
         for p_k in p:
-            _value_check((p_k >= 0) != False,
+            _value_check((p_k >= 0) != False and (p_k <= 1) != False,
                         "probability must be at least a positive symbol.")
+        _value_check(Eq(sum(p), 1) != False,
+                        "probabilities must sum to 1")
 
     @property
     def set(self):
-        return Range(0, self.n)**len(self.p)
+        i = Symbol('i', negative=False, integer=True)
+        return imageset(i, i, Interval(0, self.n))**len(self.p)
 
     def pdf(self, *x):
         n, p = self.n, self.p
         term_1 = factorial(n)/Mul.fromiter([factorial(x_k) for x_k in x])
         term_2 = Mul.fromiter([p_k**x_k for p_k, x_k in zip(p, x)])
-        return term_1*term_2
+        return Piecewise((term_1 * term_2, Eq(sum(x), n)), (0, True))
 
 def Multinomial(syms, n, *p):
     """
@@ -441,16 +444,15 @@ def Multinomial(syms, n, *p):
     >>> from sympy.stats import density
     >>> from sympy.stats.joint_rv import marginal_distribution
     >>> from sympy.stats.joint_rv_types import Multinomial
-    >>> from sympy import Symbol
     >>> from sympy import symbols
     >>> x1, x2, x3 = symbols('x1, x2, x3', nonnegative=True, integer=True)
     >>> p1, p2, p3 = symbols('p1, p2, p3', positive=True)
     >>> M = Multinomial('M', 3, p1, p2, p3)
     >>> density(M)(x1, x2, x3)
-    6*p1**x1*p2**x2*p3**x3/(factorial(x1)*factorial(x2)*factorial(x3))
+    Piecewise((6*p1**x1*p2**x2*p3**x3/(factorial(x1)*factorial(x2)*factorial(x3)),
+    Eq(x1 + x2 + x3, 3)), (0, True))
     >>> marginal_distribution(M, M[0])(x1).subs(x1, 1)
-    3*p1*p2**2*p3**2/2 + 3*p1*p2**2*p3 + 3*p1*p2**2 + 3*p1*p2*p3**2 +
-    6*p1*p2*p3 + 6*p1*p2 + 3*p1*p3**2 + 6*p1*p3 + 6*p1
+    3*p1*p2**2 + 6*p1*p2*p3 + 3*p1*p3**2
 
     References
     ==========
@@ -471,21 +473,23 @@ class NegativeMultinomialDistribution(JointDistribution):
     is_Discrete = True
 
     def check(self, k0, p):
-        _value_check(((k0 > 0) != False) and isinstance(k0, Integer),
+        _value_check(((k0 > 0) != False),
                         "number of failures must be a positve integer")
         for p_k in p:
             _value_check((p_k >= 0) != False,
                         "probability must be at least a positive symbol.")
+        _value_check((sum(p) <= 1) != False,
+                        "success probabilities must not be greater than 1.")
 
     @property
     def set(self):
-        return S.Naturals0**len(self.p)
+        return Range(0, S.Infinity)**len(self.p)
 
     def pdf(self, *k):
         k0, p = self.k0, self.p
         term_1 = (gamma(k0 + sum(k))*(1 - sum(p))**k0)/gamma(k0)
         term_2 = Mul.fromiter([pi**ki/factorial(ki) for pi, ki in zip(p, k)])
-        return term_1*term_2
+        return term_1 * term_2
 
 def NegativeMultinomial(syms, k0, *p):
     """

sympy/stats/tests/test_joint_rv.py

@@ -1,5 +1,5 @@
 from sympy import (symbols, pi, oo, S, exp, sqrt, besselk, Indexed, Rational,
-                    simplify, Piecewise, factorial, Eq, gamma)
+                    simplify, Piecewise, factorial, Eq, gamma, Sum)
 from sympy.stats import density
 from sympy.stats.joint_rv import marginal_distribution
 from sympy.stats.joint_rv_types import JointRV
@@ -93,31 +93,36 @@ def test_Multinomial():
     from sympy.stats.joint_rv_types import Multinomial
     n, x1, x2, x3, x4 = symbols('n, x1, x2, x3, x4', nonnegative=True, integer=True)
     p1, p2, p3, p4 = symbols('p1, p2, p3, p4', positive=True)
-    p1_f = symbols('p1_f', negative=True)
-    M = Multinomial('M', 3, [p1, p2, p3, p4])
-    M_c = Multinomial('C', 3, 0.5, 0.4, 0.3, 0.2)
+    p1_f, n_f = symbols('p1_f, n_f', negative=True)
+    M = Multinomial('M', n, [p1, p2, p3, p4])
+    C = Multinomial('C', n, p1, p2, p3)
     f = factorial
-    assert simplify(density(M)(x1, x2, x3, x4) -
-            S(6)*p1**x1*p2**x2*p3**x3*p4**x4/(f(x1)*f(x2)*f(x3)*f(x4))) == S(0)
-    assert marginal_distribution(M_c, M_c[0])(1).round(2) == 7.29
+    assert density(M)(x1, x2, x3, x4) == Piecewise((p1**x1*p2**x2*p3**x3*p4**x4*
+                                            f(n)/(f(x1)*f(x2)*f(x3)*f(x4)),
+                                            Eq(n, x1 + x2 + x3 + x4)), (0, True))
+    marg = Sum(Piecewise((p1**x1*p2**C[1]*p3**C[2]*factorial(n)/(factorial(x1)*
+            factorial(C[1])*factorial(C[2])), Eq(n, x1 + C[1] + C[2])), (0, True)
+            ), (C[1], 0, n), (C[2], 0, n))
+    assert str(marginal_distribution(C, C[0])(x1)) == str(marg)
     raises(ValueError, lambda: Multinomial('b1', 5, [p1, p2, p3, p1_f]))
-    raises(ValueError, lambda: Multinomial('b2', n, [p1, p2, p3, p4]))
+    raises(ValueError, lambda: Multinomial('b2', n_f, [p1, p2, p3, p4]))
+    raises(ValueError, lambda: Multinomial('b3', n, 0.5, 0.4, 0.3, 0.1))
 
 def test_NegativeMultinomial():
     from sympy.stats.joint_rv_types import NegativeMultinomial
     k0, x1, x2, x3, x4 = symbols('k0, x1, x2, x3, x4', nonnegative=True, integer=True)
     p1, p2, p3, p4 = symbols('p1, p2, p3, p4', positive=True)
     p1_f = symbols('p1_f', negative=True)
     N = NegativeMultinomial('N', 4, [p1, p2, p3, p4])
-    N_c = NegativeMultinomial('C', 4, 0.1, 0.2, 0.3)
+    C = NegativeMultinomial('C', 4, 0.1, 0.2, 0.3)
     g = gamma
     f = factorial
     assert simplify(density(N)(x1, x2, x3, x4) -
             p1**x1*p2**x2*p3**x3*p4**x4*(-p1 - p2 - p3 - p4 + 1)**4*g(x1 + x2 +
             x3 + x4 + 4)/(6*f(x1)*f(x2)*f(x3)*f(x4))) == S(0)
-    assert marginal_distribution(N_c, N_c[0])(1).evalf().round(2) == 0.33
+    assert marginal_distribution(C, C[0])(1).evalf().round(2) == 0.33
     raises(ValueError, lambda: NegativeMultinomial('b1', 5, [p1, p2, p3, p1_f]))
-    raises(ValueError, lambda: NegativeMultinomial('b2', k0, [p1, p2, p3, p4]))
+    raises(ValueError, lambda: NegativeMultinomial('b2', k0, 0.5, 0.4, 0.3, 0.4))
 
 def test_JointPSpace_margial_distribution():
     from sympy.stats.joint_rv_types import MultivariateT