from sympy import (S, Symbol, Interval, binomial, nan, exp, Or, symbols, Eq, cos, And, Tuple, integrate, oo, sin, Sum, Basic, Indexed, DiracDelta, Lambda, log, pi, FallingFactorial, Rational, Matrix) from sympy.stats import (Die, Normal, Exponential, FiniteRV, P, E, H, variance, density, given, independent, dependent, where, pspace, GaussianUnitaryEnsemble, random_symbols, sample, Geometric, factorial_moment, Binomial, Hypergeometric, DiscreteUniform, Poisson, characteristic_function, moment_generating_function, BernoulliProcess, Variance, Expectation, Probability, Covariance, covariance, cmoment, moment, median) from sympy.stats.rv import (IndependentProductPSpace, rs_swap, Density, NamedArgsMixin, RandomSymbol, sample_iter, PSpace, is_random, RandomIndexedSymbol, RandomMatrixSymbol) from sympy.testing.pytest import raises, skip, XFAIL from sympy.external import import_module from sympy.core.numbers import comp from sympy.stats.frv_types import BernoulliDistribution from sympy.core.symbol import Dummy from sympy.functions.elementary.piecewise import Piecewise def test_where(): X, Y = Die('X'), Die('Y') Z = Normal('Z', 0, 1) assert where(Z**2 <= 1).set == Interval(-1, 1) assert where(Z**2 <= 1).as_boolean() == Interval(-1, 1).as_relational(Z.symbol) assert where(And(X > Y, Y > 4)).as_boolean() == And( Eq(X.symbol, 6), Eq(Y.symbol, 5)) assert len(where(X < 3).set) == 2 assert 1 in where(X < 3).set X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) assert where(And(X**2 <= 1, X >= 0)).set == Interval(0, 1) XX = given(X, And(X**2 <= 1, X >= 0)) assert XX.pspace.domain.set == Interval(0, 1) assert XX.pspace.domain.as_boolean() == \ And(0 <= X.symbol, X.symbol**2 <= 1, -oo < X.symbol, X.symbol < oo) with raises(TypeError): XX = given(X, X + 3) def test_random_symbols(): X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) assert set(random_symbols(2*X + 1)) == {X} assert set(random_symbols(2*X + Y)) == {X, Y} assert set(random_symbols(2*X + Y.symbol)) == {X} assert set(random_symbols(2)) == set() def test_characteristic_function(): # Imports I from sympy from sympy import I X = Normal('X',0,1) Y = DiscreteUniform('Y', [1,2,7]) Z = Poisson('Z', 2) t = symbols('_t') P = Lambda(t, exp(-t**2/2)) Q = Lambda(t, exp(7*t*I)/3 + exp(2*t*I)/3 + exp(t*I)/3) R = Lambda(t, exp(2 * exp(t*I) - 2)) assert characteristic_function(X).dummy_eq(P) assert characteristic_function(Y).dummy_eq(Q) assert characteristic_function(Z).dummy_eq(R) def test_moment_generating_function(): X = Normal('X',0,1) Y = DiscreteUniform('Y', [1,2,7]) Z = Poisson('Z', 2) t = symbols('_t') P = Lambda(t, exp(t**2/2)) Q = Lambda(t, (exp(7*t)/3 + exp(2*t)/3 + exp(t)/3)) R = Lambda(t, exp(2 * exp(t) - 2)) assert moment_generating_function(X).dummy_eq(P) assert moment_generating_function(Y).dummy_eq(Q) assert moment_generating_function(Z).dummy_eq(R) def test_sample_iter(): X = Normal('X',0,1) Y = DiscreteUniform('Y', [1, 2, 7]) Z = Poisson('Z', 2) scipy = import_module('scipy') if not scipy: skip('Scipy is not installed. Abort tests') expr = X**2 + 3 iterator = sample_iter(expr) expr2 = Y**2 + 5*Y + 4 iterator2 = sample_iter(expr2) expr3 = Z**3 + 4 iterator3 = sample_iter(expr3) def is_iterator(obj): if ( hasattr(obj, '__iter__') and (hasattr(obj, 'next') or hasattr(obj, '__next__')) and callable(obj.__iter__) and obj.__iter__() is obj ): return True else: return False assert is_iterator(iterator) assert is_iterator(iterator2) assert is_iterator(iterator3) def test_pspace(): X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) x = Symbol('x') raises(ValueError, lambda: pspace(5 + 3)) raises(ValueError, lambda: pspace(x < 1)) assert pspace(X) == X.pspace assert pspace(2*X + 1) == X.pspace assert pspace(2*X + Y) == IndependentProductPSpace(Y.pspace, X.pspace) def test_rs_swap(): X = Normal('x', 0, 1) Y = Exponential('y', 1) XX = Normal('x', 0, 2) YY = Normal('y', 0, 3) expr = 2*X + Y assert expr.subs(rs_swap((X, Y), (YY, XX))) == 2*XX + YY def test_RandomSymbol(): X = Normal('x', 0, 1) Y = Normal('x', 0, 2) assert X.symbol == Y.symbol assert X != Y assert X.name == X.symbol.name X = Normal('lambda', 0, 1) # make sure we can use protected terms X = Normal('Lambda', 0, 1) # make sure we can use SymPy terms def test_RandomSymbol_diff(): X = Normal('x', 0, 1) assert (2*X).diff(X) def test_random_symbol_no_pspace(): x = RandomSymbol(Symbol('x')) assert x.pspace == PSpace() def test_overlap(): X = Normal('x', 0, 1) Y = Normal('x', 0, 2) raises(ValueError, lambda: P(X > Y)) def test_IndependentProductPSpace(): X = Normal('X', 0, 1) Y = Normal('Y', 0, 1) px = X.pspace py = Y.pspace assert pspace(X + Y) == IndependentProductPSpace(px, py) assert pspace(X + Y) == IndependentProductPSpace(py, px) def test_E(): assert E(5) == 5 def test_H(): X = Normal('X', 0, 1) D = Die('D', sides = 4) G = Geometric('G', 0.5) assert H(X, X > 0) == -log(2)/2 + S.Half + log(pi)/2 assert H(D, D > 2) == log(2) assert comp(H(G).evalf().round(2), 1.39) def test_Sample(): X = Die('X', 6) Y = Normal('Y', 0, 1) z = Symbol('z', integer=True) scipy = import_module('scipy') if not scipy: skip('Scipy is not installed. Abort tests') assert sample(X) in [1, 2, 3, 4, 5, 6] assert isinstance(sample(X + Y), float) assert P(X + Y > 0, Y < 0, numsamples=10).is_number assert E(X + Y, numsamples=10).is_number assert E(X**2 + Y, numsamples=10).is_number assert E((X + Y)**2, numsamples=10).is_number assert variance(X + Y, numsamples=10).is_number raises(TypeError, lambda: P(Y > z, numsamples=5)) assert P(sin(Y) <= 1, numsamples=10) == 1 assert P(sin(Y) <= 1, cos(Y) < 1, numsamples=10) == 1 assert all(i in range(1, 7) for i in density(X, numsamples=10)) assert all(i in range(4, 7) for i in density(X, X>3, numsamples=10)) numpy = import_module('numpy') if not numpy: skip('Numpy is not installed. Abort tests') #Test Issue #21563: Output of sample must be a float or array assert isinstance(sample(X), numpy.int64) assert isinstance(sample(Y), numpy.float64) assert isinstance(sample(X, size=2), numpy.ndarray) @XFAIL def test_samplingE(): scipy = import_module('scipy') if not scipy: skip('Scipy is not installed. Abort tests') Y = Normal('Y', 0, 1) z = Symbol('z', integer=True) assert E(Sum(1/z**Y, (z, 1, oo)), Y > 2, numsamples=3).is_number def test_given(): X = Normal('X', 0, 1) Y = Normal('Y', 0, 1) A = given(X, True) B = given(X, Y > 2) assert X == A == B def test_factorial_moment(): X = Poisson('X', 2) Y = Binomial('Y', 2, S.Half) Z = Hypergeometric('Z', 4, 2, 2) assert factorial_moment(X, 2) == 4 assert factorial_moment(Y, 2) == S.Half assert factorial_moment(Z, 2) == Rational(1, 3) x, y, z, l = symbols('x y z l') Y = Binomial('Y', 2, y) Z = Hypergeometric('Z', 10, 2, 3) assert factorial_moment(Y, l) == y**2*FallingFactorial( 2, l) + 2*y*(1 - y)*FallingFactorial(1, l) + (1 - y)**2*\ FallingFactorial(0, l) assert factorial_moment(Z, l) == 7*FallingFactorial(0, l)/\ 15 + 7*FallingFactorial(1, l)/15 + FallingFactorial(2, l)/15 def test_dependence(): X, Y = Die('X'), Die('Y') assert independent(X, 2*Y) assert not dependent(X, 2*Y) X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) assert independent(X, Y) assert dependent(X, 2*X) # Create a dependency XX, YY = given(Tuple(X, Y), Eq(X + Y, 3)) assert dependent(XX, YY) def test_dependent_finite(): X, Y = Die('X'), Die('Y') # Dependence testing requires symbolic conditions which currently break # finite random variables assert dependent(X, Y + X) XX, YY = given(Tuple(X, Y), X + Y > 5) # Create a dependency assert dependent(XX, YY) def test_normality(): X, Y = Normal('X', 0, 1), Normal('Y', 0, 1) x = Symbol('x', real=True, finite=True) z = Symbol('z', real=True, finite=True) dens = density(X - Y, Eq(X + Y, z)) assert integrate(dens(x), (x, -oo, oo)) == 1 def test_Density(): X = Die('X', 6) d = Density(X) assert d.doit() == density(X) def test_NamedArgsMixin(): class Foo(Basic, NamedArgsMixin): _argnames = 'foo', 'bar' a = Foo(1, 2) assert a.foo == 1 assert a.bar == 2 raises(AttributeError, lambda: a.baz) class Bar(Basic, NamedArgsMixin): pass raises(AttributeError, lambda: Bar(1, 2).foo) def test_density_constant(): assert density(3)(2) == 0 assert density(3)(3) == DiracDelta(0) def test_cmoment_constant(): assert variance(3) == 0 assert cmoment(3, 3) == 0 assert cmoment(3, 4) == 0 x = Symbol('x') assert variance(x) == 0 assert cmoment(x, 15) == 0 assert cmoment(x, 0) == 1 def test_moment_constant(): assert moment(3, 0) == 1 assert moment(3, 1) == 3 assert moment(3, 2) == 9 x = Symbol('x') assert moment(x, 2) == x**2 def test_median_constant(): assert median(3) == 3 x = Symbol('x') assert median(x) == x def test_real(): x = Normal('x', 0, 1) assert x.is_real def test_issue_10052(): X = Exponential('X', 3) assert P(X < oo) == 1 assert P(X > oo) == 0 assert P(X < 2, X > oo) == 0 assert P(X < oo, X > oo) == 0 assert P(X < oo, X > 2) == 1 assert P(X < 3, X == 2) == 0 raises(ValueError, lambda: P(1)) raises(ValueError, lambda: P(X < 1, 2)) def test_issue_11934(): density = {0: .5, 1: .5} X = FiniteRV('X', density) assert E(X) == 0.5 assert P( X>= 2) == 0 def test_issue_8129(): X = Exponential('X', 4) assert P(X >= X) == 1 assert P(X > X) == 0 assert P(X > X+1) == 0 def test_issue_12237(): X = Normal('X', 0, 1) Y = Normal('Y', 0, 1) U = P(X > 0, X) V = P(Y < 0, X) W = P(X + Y > 0, X) assert W == P(X + Y > 0, X) assert U == BernoulliDistribution(S.Half, S.Zero, S.One) assert V == S.Half def test_is_random(): X = Normal('X', 0, 1) Y = Normal('Y', 0, 1) a, b = symbols('a, b') G = GaussianUnitaryEnsemble('U', 2) B = BernoulliProcess('B', 0.9) assert not is_random(a) assert not is_random(a + b) assert not is_random(a * b) assert not is_random(Matrix([a**2, b**2])) assert is_random(X) assert is_random(X**2 + Y) assert is_random(Y + b**2) assert is_random(Y > 5) assert is_random(B[3] < 1) assert is_random(G) assert is_random(X * Y * B[1]) assert is_random(Matrix([[X, B[2]], [G, Y]])) assert is_random(Eq(X, 4)) def test_issue_12283(): x = symbols('x') X = RandomSymbol(x) Y = RandomSymbol('Y') Z = RandomMatrixSymbol('Z', 2, 1) W = RandomMatrixSymbol('W', 2, 1) RI = RandomIndexedSymbol(Indexed('RI', 3)) assert pspace(Z) == PSpace() assert pspace(RI) == PSpace() assert pspace(X) == PSpace() assert E(X) == Expectation(X) assert P(Y > 3) == Probability(Y > 3) assert variance(X) == Variance(X) assert variance(RI) == Variance(RI) assert covariance(X, Y) == Covariance(X, Y) assert covariance(W, Z) == Covariance(W, Z) def test_issue_6810(): X = Die('X', 6) Y = Normal('Y', 0, 1) assert P(Eq(X, 2)) == S(1)/6 assert P(Eq(Y, 0)) == 0 assert P(Or(X > 2, X < 3)) == 1 assert P(And(X > 3, X > 2)) == S(1)/2 def test_issue_20286(): n, p = symbols('n p') B = Binomial('B', n, p) k = Dummy('k', integer = True) eq = Sum(Piecewise((-p**k*(1 - p)**(-k + n)*log(p**k*(1 - p)**(-k + n)*binomial(n, k))*binomial(n, k), (k >= 0) & (k <= n)), (nan, True)), (k, 0, n)) assert eq.dummy_eq(H(B))