Spaces:
Sleeping
Sleeping
| from sympy.concrete.summations import Sum | |
| from sympy.core.numbers import (oo, pi) | |
| from sympy.core.relational import Eq | |
| from sympy.core.singleton import S | |
| from sympy.core.symbol import symbols | |
| from sympy.functions.combinatorial.factorials import factorial | |
| from sympy.functions.elementary.exponential import exp | |
| from sympy.functions.elementary.miscellaneous import sqrt | |
| from sympy.functions.elementary.piecewise import Piecewise | |
| from sympy.functions.special.beta_functions import beta | |
| from sympy.functions.special.error_functions import erf | |
| from sympy.functions.special.gamma_functions import gamma | |
| from sympy.integrals.integrals import Integral | |
| from sympy.sets.sets import Interval | |
| from sympy.stats import (Normal, P, E, density, Gamma, Poisson, Rayleigh, | |
| variance, Bernoulli, Beta, Uniform, cdf) | |
| from sympy.stats.compound_rv import CompoundDistribution, CompoundPSpace | |
| from sympy.stats.crv_types import NormalDistribution | |
| from sympy.stats.drv_types import PoissonDistribution | |
| from sympy.stats.frv_types import BernoulliDistribution | |
| from sympy.testing.pytest import raises, ignore_warnings | |
| from sympy.stats.joint_rv_types import MultivariateNormalDistribution | |
| from sympy.abc import x | |
| # helpers for testing troublesome unevaluated expressions | |
| flat = lambda s: ''.join(str(s).split()) | |
| streq = lambda *a: len(set(map(flat, a))) == 1 | |
| assert streq(x, x) | |
| assert streq(x, 'x') | |
| assert not streq(x, x + 1) | |
| def test_normal_CompoundDist(): | |
| X = Normal('X', 1, 2) | |
| Y = Normal('X', X, 4) | |
| assert density(Y)(x).simplify() == sqrt(10)*exp(-x**2/40 + x/20 - S(1)/40)/(20*sqrt(pi)) | |
| assert E(Y) == 1 # it is always equal to mean of X | |
| assert P(Y > 1) == S(1)/2 # as 1 is the mean | |
| assert P(Y > 5).simplify() == S(1)/2 - erf(sqrt(10)/5)/2 | |
| assert variance(Y) == variance(X) + 4**2 # 2**2 + 4**2 | |
| # https://math.stackexchange.com/questions/1484451/ | |
| # (Contains proof of E and variance computation) | |
| def test_poisson_CompoundDist(): | |
| k, t, y = symbols('k t y', positive=True, real=True) | |
| G = Gamma('G', k, t) | |
| D = Poisson('P', G) | |
| assert density(D)(y).simplify() == t**y*(t + 1)**(-k - y)*gamma(k + y)/(gamma(k)*gamma(y + 1)) | |
| # https://en.wikipedia.org/wiki/Negative_binomial_distribution#Gamma%E2%80%93Poisson_mixture | |
| assert E(D).simplify() == k*t # mean of NegativeBinomialDistribution | |
| def test_bernoulli_CompoundDist(): | |
| X = Beta('X', 1, 2) | |
| Y = Bernoulli('Y', X) | |
| assert density(Y).dict == {0: S(2)/3, 1: S(1)/3} | |
| assert E(Y) == P(Eq(Y, 1)) == S(1)/3 | |
| assert variance(Y) == S(2)/9 | |
| assert cdf(Y) == {0: S(2)/3, 1: 1} | |
| # test issue 8128 | |
| a = Bernoulli('a', S(1)/2) | |
| b = Bernoulli('b', a) | |
| assert density(b).dict == {0: S(1)/2, 1: S(1)/2} | |
| assert P(b > 0.5) == S(1)/2 | |
| X = Uniform('X', 0, 1) | |
| Y = Bernoulli('Y', X) | |
| assert E(Y) == S(1)/2 | |
| assert P(Eq(Y, 1)) == E(Y) | |
| def test_unevaluated_CompoundDist(): | |
| # these tests need to be removed once they work with evaluation as they are currently not | |
| # evaluated completely in sympy. | |
| R = Rayleigh('R', 4) | |
| X = Normal('X', 3, R) | |
| ans = ''' | |
| Piecewise(((-sqrt(pi)*sinh(x/4 - 3/4) + sqrt(pi)*cosh(x/4 - 3/4))/( | |
| 8*sqrt(pi)), Abs(arg(x - 3)) <= pi/4), (Integral(sqrt(2)*exp(-(x - 3) | |
| **2/(2*R**2))*exp(-R**2/32)/(32*sqrt(pi)), (R, 0, oo)), True))''' | |
| assert streq(density(X)(x), ans) | |
| expre = ''' | |
| Integral(X*Integral(sqrt(2)*exp(-(X-3)**2/(2*R**2))*exp(-R**2/32)/(32* | |
| sqrt(pi)),(R,0,oo)),(X,-oo,oo))''' | |
| with ignore_warnings(UserWarning): ### TODO: Restore tests once warnings are removed | |
| assert streq(E(X, evaluate=False).rewrite(Integral), expre) | |
| X = Poisson('X', 1) | |
| Y = Poisson('Y', X) | |
| Z = Poisson('Z', Y) | |
| exprd = Sum(exp(-Y)*Y**x*Sum(exp(-1)*exp(-X)*X**Y/(factorial(X)*factorial(Y) | |
| ), (X, 0, oo))/factorial(x), (Y, 0, oo)) | |
| assert density(Z)(x) == exprd | |
| N = Normal('N', 1, 2) | |
| M = Normal('M', 3, 4) | |
| D = Normal('D', M, N) | |
| exprd = ''' | |
| Integral(sqrt(2)*exp(-(N-1)**2/8)*Integral(exp(-(x-M)**2/(2*N**2))*exp | |
| (-(M-3)**2/32)/(8*pi*N),(M,-oo,oo))/(4*sqrt(pi)),(N,-oo,oo))''' | |
| assert streq(density(D, evaluate=False)(x), exprd) | |
| def test_Compound_Distribution(): | |
| X = Normal('X', 2, 4) | |
| N = NormalDistribution(X, 4) | |
| C = CompoundDistribution(N) | |
| assert C.is_Continuous | |
| assert C.set == Interval(-oo, oo) | |
| assert C.pdf(x, evaluate=True).simplify() == exp(-x**2/64 + x/16 - S(1)/16)/(8*sqrt(pi)) | |
| assert not isinstance(CompoundDistribution(NormalDistribution(2, 3)), | |
| CompoundDistribution) | |
| M = MultivariateNormalDistribution([1, 2], [[2, 1], [1, 2]]) | |
| raises(NotImplementedError, lambda: CompoundDistribution(M)) | |
| X = Beta('X', 2, 4) | |
| B = BernoulliDistribution(X, 1, 0) | |
| C = CompoundDistribution(B) | |
| assert C.is_Finite | |
| assert C.set == {0, 1} | |
| y = symbols('y', negative=False, integer=True) | |
| assert C.pdf(y, evaluate=True) == Piecewise((S(1)/(30*beta(2, 4)), Eq(y, 0)), | |
| (S(1)/(60*beta(2, 4)), Eq(y, 1)), (0, True)) | |
| k, t, z = symbols('k t z', positive=True, real=True) | |
| G = Gamma('G', k, t) | |
| X = PoissonDistribution(G) | |
| C = CompoundDistribution(X) | |
| assert C.is_Discrete | |
| assert C.set == S.Naturals0 | |
| assert C.pdf(z, evaluate=True).simplify() == t**z*(t + 1)**(-k - z)*gamma(k \ | |
| + z)/(gamma(k)*gamma(z + 1)) | |
| def test_compound_pspace(): | |
| X = Normal('X', 2, 4) | |
| Y = Normal('Y', 3, 6) | |
| assert not isinstance(Y.pspace, CompoundPSpace) | |
| N = NormalDistribution(1, 2) | |
| D = PoissonDistribution(3) | |
| B = BernoulliDistribution(0.2, 1, 0) | |
| pspace1 = CompoundPSpace('N', N) | |
| pspace2 = CompoundPSpace('D', D) | |
| pspace3 = CompoundPSpace('B', B) | |
| assert not isinstance(pspace1, CompoundPSpace) | |
| assert not isinstance(pspace2, CompoundPSpace) | |
| assert not isinstance(pspace3, CompoundPSpace) | |
| M = MultivariateNormalDistribution([1, 2], [[2, 1], [1, 2]]) | |
| raises(ValueError, lambda: CompoundPSpace('M', M)) | |
| Y = Normal('Y', X, 6) | |
| assert isinstance(Y.pspace, CompoundPSpace) | |
| assert Y.pspace.distribution == CompoundDistribution(NormalDistribution(X, 6)) | |
| assert Y.pspace.domain.set == Interval(-oo, oo) | |