510 lines
20 KiB
Python
510 lines
20 KiB
Python
|
from sympy.concrete.summations import Sum
|
||
|
|
from sympy.core.containers import (Dict, Tuple)
|
||
|
|
from sympy.core.function import Function
|
||
|
|
from sympy.core.numbers import (I, Rational, nan)
|
||
|
|
from sympy.core.relational import Eq
|
||
|
|
from sympy.core.singleton import S
|
||
|
|
from sympy.core.symbol import (Dummy, Symbol, symbols)
|
||
|
|
from sympy.core.sympify import sympify
|
||
|
|
from sympy.functions.combinatorial.factorials import binomial
|
||
|
|
from sympy.functions.combinatorial.numbers import harmonic
|
||
|
|
from sympy.functions.elementary.exponential import exp
|
||
|
|
from sympy.functions.elementary.miscellaneous import sqrt
|
||
|
|
from sympy.functions.elementary.piecewise import Piecewise
|
||
|
|
from sympy.functions.elementary.trigonometric import cos
|
||
|
|
from sympy.functions.special.beta_functions import beta
|
||
|
|
from sympy.logic.boolalg import (And, Or)
|
||
|
|
from sympy.polys.polytools import cancel
|
||
|
|
from sympy.sets.sets import FiniteSet
|
||
|
|
from sympy.simplify.simplify import simplify
|
||
|
|
from sympy.matrices import Matrix
|
||
|
|
from sympy.stats import (DiscreteUniform, Die, Bernoulli, Coin, Binomial, BetaBinomial,
|
||
|
|
Hypergeometric, Rademacher, IdealSoliton, RobustSoliton, P, E, variance,
|
||
|
|
covariance, skewness, density, where, FiniteRV, pspace, cdf,
|
||
|
|
correlation, moment, cmoment, smoment, characteristic_function,
|
||
|
|
moment_generating_function, quantile, kurtosis, median, coskewness)
|
||
|
|
from sympy.stats.frv_types import DieDistribution, BinomialDistribution, \
|
||
|
|
HypergeometricDistribution
|
||
|
|
from sympy.stats.rv import Density
|
||
|
|
from sympy.testing.pytest import raises
|
||
|
|
|
||
|
|
|
||
|
|
def BayesTest(A, B):
|
||
|
|
assert P(A, B) == P(And(A, B)) / P(B)
|
||
|
|
assert P(A, B) == P(B, A) * P(A) / P(B)
|
||
|
|
|
||
|
|
|
||
|
|
def test_discreteuniform():
|
||
|
|
# Symbolic
|
||
|
|
a, b, c, t = symbols('a b c t')
|
||
|
|
X = DiscreteUniform('X', [a, b, c])
|
||
|
|
|
||
|
|
assert E(X) == (a + b + c)/3
|
||
|
|
assert simplify(variance(X)
|
||
|
|
- ((a**2 + b**2 + c**2)/3 - (a/3 + b/3 + c/3)**2)) == 0
|
||
|
|
assert P(Eq(X, a)) == P(Eq(X, b)) == P(Eq(X, c)) == S('1/3')
|
||
|
|
|
||
|
|
Y = DiscreteUniform('Y', range(-5, 5))
|
||
|
|
|
||
|
|
# Numeric
|
||
|
|
assert E(Y) == S('-1/2')
|
||
|
|
assert variance(Y) == S('33/4')
|
||
|
|
assert median(Y) == FiniteSet(-1, 0)
|
||
|
|
|
||
|
|
for x in range(-5, 5):
|
||
|
|
assert P(Eq(Y, x)) == S('1/10')
|
||
|
|
assert P(Y <= x) == S(x + 6)/10
|
||
|
|
assert P(Y >= x) == S(5 - x)/10
|
||
|
|
|
||
|
|
assert dict(density(Die('D', 6)).items()) == \
|
||
|
|
dict(density(DiscreteUniform('U', range(1, 7))).items())
|
||
|
|
|
||
|
|
assert characteristic_function(X)(t) == exp(I*a*t)/3 + exp(I*b*t)/3 + exp(I*c*t)/3
|
||
|
|
assert moment_generating_function(X)(t) == exp(a*t)/3 + exp(b*t)/3 + exp(c*t)/3
|
||
|
|
# issue 18611
|
||
|
|
raises(ValueError, lambda: DiscreteUniform('Z', [a, a, a, b, b, c]))
|
||
|
|
|
||
|
|
def test_dice():
|
||
|
|
# TODO: Make iid method!
|
||
|
|
X, Y, Z = Die('X', 6), Die('Y', 6), Die('Z', 6)
|
||
|
|
a, b, t, p = symbols('a b t p')
|
||
|
|
|
||
|
|
assert E(X) == 3 + S.Half
|
||
|
|
assert variance(X) == Rational(35, 12)
|
||
|
|
assert E(X + Y) == 7
|
||
|
|
assert E(X + X) == 7
|
||
|
|
assert E(a*X + b) == a*E(X) + b
|
||
|
|
assert variance(X + Y) == variance(X) + variance(Y) == cmoment(X + Y, 2)
|
||
|
|
assert variance(X + X) == 4 * variance(X) == cmoment(X + X, 2)
|
||
|
|
assert cmoment(X, 0) == 1
|
||
|
|
assert cmoment(4*X, 3) == 64*cmoment(X, 3)
|
||
|
|
assert covariance(X, Y) is S.Zero
|
||
|
|
assert covariance(X, X + Y) == variance(X)
|
||
|
|
assert density(Eq(cos(X*S.Pi), 1))[True] == S.Half
|
||
|
|
assert correlation(X, Y) == 0
|
||
|
|
assert correlation(X, Y) == correlation(Y, X)
|
||
|
|
assert smoment(X + Y, 3) == skewness(X + Y)
|
||
|
|
assert smoment(X + Y, 4) == kurtosis(X + Y)
|
||
|
|
assert smoment(X, 0) == 1
|
||
|
|
assert P(X > 3) == S.Half
|
||
|
|
assert P(2*X > 6) == S.Half
|
||
|
|
assert P(X > Y) == Rational(5, 12)
|
||
|
|
assert P(Eq(X, Y)) == P(Eq(X, 1))
|
||
|
|
|
||
|
|
assert E(X, X > 3) == 5 == moment(X, 1, 0, X > 3)
|
||
|
|
assert E(X, Y > 3) == E(X) == moment(X, 1, 0, Y > 3)
|
||
|
|
assert E(X + Y, Eq(X, Y)) == E(2*X)
|
||
|
|
assert moment(X, 0) == 1
|
||
|
|
assert moment(5*X, 2) == 25*moment(X, 2)
|
||
|
|
assert quantile(X)(p) == Piecewise((nan, (p > 1) | (p < 0)),\
|
||
|
|
(S.One, p <= Rational(1, 6)), (S(2), p <= Rational(1, 3)), (S(3), p <= S.Half),\
|
||
|
|
(S(4), p <= Rational(2, 3)), (S(5), p <= Rational(5, 6)), (S(6), p <= 1))
|
||
|
|
|
||
|
|
assert P(X > 3, X > 3) is S.One
|
||
|
|
assert P(X > Y, Eq(Y, 6)) is S.Zero
|
||
|
|
assert P(Eq(X + Y, 12)) == Rational(1, 36)
|
||
|
|
assert P(Eq(X + Y, 12), Eq(X, 6)) == Rational(1, 6)
|
||
|
|
|
||
|
|
assert density(X + Y) == density(Y + Z) != density(X + X)
|
||
|
|
d = density(2*X + Y**Z)
|
||
|
|
assert d[S(22)] == Rational(1, 108) and d[S(4100)] == Rational(1, 216) and S(3130) not in d
|
||
|
|
|
||
|
|
assert pspace(X).domain.as_boolean() == Or(
|
||
|
|
*[Eq(X.symbol, i) for i in [1, 2, 3, 4, 5, 6]])
|
||
|
|
|
||
|
|
assert where(X > 3).set == FiniteSet(4, 5, 6)
|
||
|
|
|
||
|
|
assert characteristic_function(X)(t) == exp(6*I*t)/6 + exp(5*I*t)/6 + exp(4*I*t)/6 + exp(3*I*t)/6 + exp(2*I*t)/6 + exp(I*t)/6
|
||
|
|
assert moment_generating_function(X)(t) == exp(6*t)/6 + exp(5*t)/6 + exp(4*t)/6 + exp(3*t)/6 + exp(2*t)/6 + exp(t)/6
|
||
|
|
assert median(X) == FiniteSet(3, 4)
|
||
|
|
D = Die('D', 7)
|
||
|
|
assert median(D) == FiniteSet(4)
|
||
|
|
# Bayes test for die
|
||
|
|
BayesTest(X > 3, X + Y < 5)
|
||
|
|
BayesTest(Eq(X - Y, Z), Z > Y)
|
||
|
|
BayesTest(X > 3, X > 2)
|
||
|
|
|
||
|
|
# arg test for die
|
||
|
|
raises(ValueError, lambda: Die('X', -1)) # issue 8105: negative sides.
|
||
|
|
raises(ValueError, lambda: Die('X', 0))
|
||
|
|
raises(ValueError, lambda: Die('X', 1.5)) # issue 8103: non integer sides.
|
||
|
|
|
||
|
|
# symbolic test for die
|
||
|
|
n, k = symbols('n, k', positive=True)
|
||
|
|
D = Die('D', n)
|
||
|
|
dens = density(D).dict
|
||
|
|
assert dens == Density(DieDistribution(n))
|
||
|
|
assert set(dens.subs(n, 4).doit().keys()) == {1, 2, 3, 4}
|
||
|
|
assert set(dens.subs(n, 4).doit().values()) == {Rational(1, 4)}
|
||
|
|
k = Dummy('k', integer=True)
|
||
|
|
assert E(D).dummy_eq(
|
||
|
|
Sum(Piecewise((k/n, k <= n), (0, True)), (k, 1, n)))
|
||
|
|
assert variance(D).subs(n, 6).doit() == Rational(35, 12)
|
||
|
|
|
||
|
|
ki = Dummy('ki')
|
||
|
|
cumuf = cdf(D)(k)
|
||
|
|
assert cumuf.dummy_eq(
|
||
|
|
Sum(Piecewise((1/n, (ki >= 1) & (ki <= n)), (0, True)), (ki, 1, k)))
|
||
|
|
assert cumuf.subs({n: 6, k: 2}).doit() == Rational(1, 3)
|
||
|
|
|
||
|
|
t = Dummy('t')
|
||
|
|
cf = characteristic_function(D)(t)
|
||
|
|
assert cf.dummy_eq(
|
||
|
|
Sum(Piecewise((exp(ki*I*t)/n, (ki >= 1) & (ki <= n)), (0, True)), (ki, 1, n)))
|
||
|
|
assert cf.subs(n, 3).doit() == exp(3*I*t)/3 + exp(2*I*t)/3 + exp(I*t)/3
|
||
|
|
mgf = moment_generating_function(D)(t)
|
||
|
|
assert mgf.dummy_eq(
|
||
|
|
Sum(Piecewise((exp(ki*t)/n, (ki >= 1) & (ki <= n)), (0, True)), (ki, 1, n)))
|
||
|
|
assert mgf.subs(n, 3).doit() == exp(3*t)/3 + exp(2*t)/3 + exp(t)/3
|
||
|
|
|
||
|
|
def test_given():
|
||
|
|
X = Die('X', 6)
|
||
|
|
assert density(X, X > 5) == {S(6): S.One}
|
||
|
|
assert where(X > 2, X > 5).as_boolean() == Eq(X.symbol, 6)
|
||
|
|
|
||
|
|
|
||
|
|
def test_domains():
|
||
|
|
X, Y = Die('x', 6), Die('y', 6)
|
||
|
|
x, y = X.symbol, Y.symbol
|
||
|
|
# Domains
|
||
|
|
d = where(X > Y)
|
||
|
|
assert d.condition == (x > y)
|
||
|
|
d = where(And(X > Y, Y > 3))
|
||
|
|
assert d.as_boolean() == Or(And(Eq(x, 5), Eq(y, 4)), And(Eq(x, 6),
|
||
|
|
Eq(y, 5)), And(Eq(x, 6), Eq(y, 4)))
|
||
|
|
assert len(d.elements) == 3
|
||
|
|
|
||
|
|
assert len(pspace(X + Y).domain.elements) == 36
|
||
|
|
|
||
|
|
Z = Die('x', 4)
|
||
|
|
|
||
|
|
raises(ValueError, lambda: P(X > Z)) # Two domains with same internal symbol
|
||
|
|
|
||
|
|
assert pspace(X + Y).domain.set == FiniteSet(1, 2, 3, 4, 5, 6)**2
|
||
|
|
|
||
|
|
assert where(X > 3).set == FiniteSet(4, 5, 6)
|
||
|
|
assert X.pspace.domain.dict == FiniteSet(
|
||
|
|
*[Dict({X.symbol: i}) for i in range(1, 7)])
|
||
|
|
|
||
|
|
assert where(X > Y).dict == FiniteSet(*[Dict({X.symbol: i, Y.symbol: j})
|
||
|
|
for i in range(1, 7) for j in range(1, 7) if i > j])
|
||
|
|
|
||
|
|
def test_bernoulli():
|
||
|
|
p, a, b, t = symbols('p a b t')
|
||
|
|
X = Bernoulli('B', p, a, b)
|
||
|
|
|
||
|
|
assert E(X) == a*p + b*(-p + 1)
|
||
|
|
assert density(X)[a] == p
|
||
|
|
assert density(X)[b] == 1 - p
|
||
|
|
assert characteristic_function(X)(t) == p * exp(I * a * t) + (-p + 1) * exp(I * b * t)
|
||
|
|
assert moment_generating_function(X)(t) == p * exp(a * t) + (-p + 1) * exp(b * t)
|
||
|
|
|
||
|
|
X = Bernoulli('B', p, 1, 0)
|
||
|
|
z = Symbol("z")
|
||
|
|
|
||
|
|
assert E(X) == p
|
||
|
|
assert simplify(variance(X)) == p*(1 - p)
|
||
|
|
assert E(a*X + b) == a*E(X) + b
|
||
|
|
assert simplify(variance(a*X + b)) == simplify(a**2 * variance(X))
|
||
|
|
assert quantile(X)(z) == Piecewise((nan, (z > 1) | (z < 0)), (0, z <= 1 - p), (1, z <= 1))
|
||
|
|
Y = Bernoulli('Y', Rational(1, 2))
|
||
|
|
assert median(Y) == FiniteSet(0, 1)
|
||
|
|
Z = Bernoulli('Z', Rational(2, 3))
|
||
|
|
assert median(Z) == FiniteSet(1)
|
||
|
|
raises(ValueError, lambda: Bernoulli('B', 1.5))
|
||
|
|
raises(ValueError, lambda: Bernoulli('B', -0.5))
|
||
|
|
|
||
|
|
#issue 8248
|
||
|
|
assert X.pspace.compute_expectation(1) == 1
|
||
|
|
|
||
|
|
p = Rational(1, 5)
|
||
|
|
X = Binomial('X', 5, p)
|
||
|
|
Y = Binomial('Y', 7, 2*p)
|
||
|
|
Z = Binomial('Z', 9, 3*p)
|
||
|
|
assert coskewness(Y + Z, X + Y, X + Z).simplify() == 0
|
||
|
|
assert coskewness(Y + 2*X + Z, X + 2*Y + Z, X + 2*Z + Y).simplify() == \
|
||
|
|
sqrt(1529)*Rational(12, 16819)
|
||
|
|
assert coskewness(Y + 2*X + Z, X + 2*Y + Z, X + 2*Z + Y, X < 2).simplify() \
|
||
|
|
== -sqrt(357451121)*Rational(2812, 4646864573)
|
||
|
|
|
||
|
|
def test_cdf():
|
||
|
|
D = Die('D', 6)
|
||
|
|
o = S.One
|
||
|
|
|
||
|
|
assert cdf(
|
||
|
|
D) == sympify({1: o/6, 2: o/3, 3: o/2, 4: 2*o/3, 5: 5*o/6, 6: o})
|
||
|
|
|
||
|
|
|
||
|
|
def test_coins():
|
||
|
|
C, D = Coin('C'), Coin('D')
|
||
|
|
H, T = symbols('H, T')
|
||
|
|
assert P(Eq(C, D)) == S.Half
|
||
|
|
assert density(Tuple(C, D)) == {(H, H): Rational(1, 4), (H, T): Rational(1, 4),
|
||
|
|
(T, H): Rational(1, 4), (T, T): Rational(1, 4)}
|
||
|
|
assert dict(density(C).items()) == {H: S.Half, T: S.Half}
|
||
|
|
|
||
|
|
F = Coin('F', Rational(1, 10))
|
||
|
|
assert P(Eq(F, H)) == Rational(1, 10)
|
||
|
|
|
||
|
|
d = pspace(C).domain
|
||
|
|
|
||
|
|
assert d.as_boolean() == Or(Eq(C.symbol, H), Eq(C.symbol, T))
|
||
|
|
|
||
|
|
raises(ValueError, lambda: P(C > D)) # Can't intelligently compare H to T
|
||
|
|
|
||
|
|
def test_binomial_verify_parameters():
|
||
|
|
raises(ValueError, lambda: Binomial('b', .2, .5))
|
||
|
|
raises(ValueError, lambda: Binomial('b', 3, 1.5))
|
||
|
|
|
||
|
|
def test_binomial_numeric():
|
||
|
|
nvals = range(5)
|
||
|
|
pvals = [0, Rational(1, 4), S.Half, Rational(3, 4), 1]
|
||
|
|
|
||
|
|
for n in nvals:
|
||
|
|
for p in pvals:
|
||
|
|
X = Binomial('X', n, p)
|
||
|
|
assert E(X) == n*p
|
||
|
|
assert variance(X) == n*p*(1 - p)
|
||
|
|
if n > 0 and 0 < p < 1:
|
||
|
|
assert skewness(X) == (1 - 2*p)/sqrt(n*p*(1 - p))
|
||
|
|
assert kurtosis(X) == 3 + (1 - 6*p*(1 - p))/(n*p*(1 - p))
|
||
|
|
for k in range(n + 1):
|
||
|
|
assert P(Eq(X, k)) == binomial(n, k)*p**k*(1 - p)**(n - k)
|
||
|
|
|
||
|
|
def test_binomial_quantile():
|
||
|
|
X = Binomial('X', 50, S.Half)
|
||
|
|
assert quantile(X)(0.95) == S(31)
|
||
|
|
assert median(X) == FiniteSet(25)
|
||
|
|
|
||
|
|
X = Binomial('X', 5, S.Half)
|
||
|
|
p = Symbol("p", positive=True)
|
||
|
|
assert quantile(X)(p) == Piecewise((nan, p > S.One), (S.Zero, p <= Rational(1, 32)),\
|
||
|
|
(S.One, p <= Rational(3, 16)), (S(2), p <= S.Half), (S(3), p <= Rational(13, 16)),\
|
||
|
|
(S(4), p <= Rational(31, 32)), (S(5), p <= S.One))
|
||
|
|
assert median(X) == FiniteSet(2, 3)
|
||
|
|
|
||
|
|
|
||
|
|
def test_binomial_symbolic():
|
||
|
|
n = 2
|
||
|
|
p = symbols('p', positive=True)
|
||
|
|
X = Binomial('X', n, p)
|
||
|
|
t = Symbol('t')
|
||
|
|
|
||
|
|
assert simplify(E(X)) == n*p == simplify(moment(X, 1))
|
||
|
|
assert simplify(variance(X)) == n*p*(1 - p) == simplify(cmoment(X, 2))
|
||
|
|
assert cancel(skewness(X) - (1 - 2*p)/sqrt(n*p*(1 - p))) == 0
|
||
|
|
assert cancel((kurtosis(X)) - (3 + (1 - 6*p*(1 - p))/(n*p*(1 - p)))) == 0
|
||
|
|
assert characteristic_function(X)(t) == p ** 2 * exp(2 * I * t) + 2 * p * (-p + 1) * exp(I * t) + (-p + 1) ** 2
|
||
|
|
assert moment_generating_function(X)(t) == p ** 2 * exp(2 * t) + 2 * p * (-p + 1) * exp(t) + (-p + 1) ** 2
|
||
|
|
|
||
|
|
# Test ability to change success/failure winnings
|
||
|
|
H, T = symbols('H T')
|
||
|
|
Y = Binomial('Y', n, p, succ=H, fail=T)
|
||
|
|
assert simplify(E(Y) - (n*(H*p + T*(1 - p)))) == 0
|
||
|
|
|
||
|
|
# test symbolic dimensions
|
||
|
|
n = symbols('n')
|
||
|
|
B = Binomial('B', n, p)
|
||
|
|
raises(NotImplementedError, lambda: P(B > 2))
|
||
|
|
assert density(B).dict == Density(BinomialDistribution(n, p, 1, 0))
|
||
|
|
assert set(density(B).dict.subs(n, 4).doit().keys()) == \
|
||
|
|
{S.Zero, S.One, S(2), S(3), S(4)}
|
||
|
|
assert set(density(B).dict.subs(n, 4).doit().values()) == \
|
||
|
|
{(1 - p)**4, 4*p*(1 - p)**3, 6*p**2*(1 - p)**2, 4*p**3*(1 - p), p**4}
|
||
|
|
k = Dummy('k', integer=True)
|
||
|
|
assert E(B > 2).dummy_eq(
|
||
|
|
Sum(Piecewise((k*p**k*(1 - p)**(-k + n)*binomial(n, k), (k >= 0)
|
||
|
|
& (k <= n) & (k > 2)), (0, True)), (k, 0, n)))
|
||
|
|
|
||
|
|
def test_beta_binomial():
|
||
|
|
# verify parameters
|
||
|
|
raises(ValueError, lambda: BetaBinomial('b', .2, 1, 2))
|
||
|
|
raises(ValueError, lambda: BetaBinomial('b', 2, -1, 2))
|
||
|
|
raises(ValueError, lambda: BetaBinomial('b', 2, 1, -2))
|
||
|
|
assert BetaBinomial('b', 2, 1, 1)
|
||
|
|
|
||
|
|
# test numeric values
|
||
|
|
nvals = range(1,5)
|
||
|
|
alphavals = [Rational(1, 4), S.Half, Rational(3, 4), 1, 10]
|
||
|
|
betavals = [Rational(1, 4), S.Half, Rational(3, 4), 1, 10]
|
||
|
|
|
||
|
|
for n in nvals:
|
||
|
|
for a in alphavals:
|
||
|
|
for b in betavals:
|
||
|
|
X = BetaBinomial('X', n, a, b)
|
||
|
|
assert E(X) == moment(X, 1)
|
||
|
|
assert variance(X) == cmoment(X, 2)
|
||
|
|
|
||
|
|
# test symbolic
|
||
|
|
n, a, b = symbols('a b n')
|
||
|
|
assert BetaBinomial('x', n, a, b)
|
||
|
|
n = 2 # Because we're using for loops, can't do symbolic n
|
||
|
|
a, b = symbols('a b', positive=True)
|
||
|
|
X = BetaBinomial('X', n, a, b)
|
||
|
|
t = Symbol('t')
|
||
|
|
|
||
|
|
assert E(X).expand() == moment(X, 1).expand()
|
||
|
|
assert variance(X).expand() == cmoment(X, 2).expand()
|
||
|
|
assert skewness(X) == smoment(X, 3)
|
||
|
|
assert characteristic_function(X)(t) == exp(2*I*t)*beta(a + 2, b)/beta(a, b) +\
|
||
|
|
2*exp(I*t)*beta(a + 1, b + 1)/beta(a, b) + beta(a, b + 2)/beta(a, b)
|
||
|
|
assert moment_generating_function(X)(t) == exp(2*t)*beta(a + 2, b)/beta(a, b) +\
|
||
|
|
2*exp(t)*beta(a + 1, b + 1)/beta(a, b) + beta(a, b + 2)/beta(a, b)
|
||
|
|
|
||
|
|
def test_hypergeometric_numeric():
|
||
|
|
for N in range(1, 5):
|
||
|
|
for m in range(0, N + 1):
|
||
|
|
for n in range(1, N + 1):
|
||
|
|
X = Hypergeometric('X', N, m, n)
|
||
|
|
N, m, n = map(sympify, (N, m, n))
|
||
|
|
assert sum(density(X).values()) == 1
|
||
|
|
assert E(X) == n * m / N
|
||
|
|
if N > 1:
|
||
|
|
assert variance(X) == n*(m/N)*(N - m)/N*(N - n)/(N - 1)
|
||
|
|
# Only test for skewness when defined
|
||
|
|
if N > 2 and 0 < m < N and n < N:
|
||
|
|
assert skewness(X) == simplify((N - 2*m)*sqrt(N - 1)*(N - 2*n)
|
||
|
|
/ (sqrt(n*m*(N - m)*(N - n))*(N - 2)))
|
||
|
|
|
||
|
|
def test_hypergeometric_symbolic():
|
||
|
|
N, m, n = symbols('N, m, n')
|
||
|
|
H = Hypergeometric('H', N, m, n)
|
||
|
|
dens = density(H).dict
|
||
|
|
expec = E(H > 2)
|
||
|
|
assert dens == Density(HypergeometricDistribution(N, m, n))
|
||
|
|
assert dens.subs(N, 5).doit() == Density(HypergeometricDistribution(5, m, n))
|
||
|
|
assert set(dens.subs({N: 3, m: 2, n: 1}).doit().keys()) == {S.Zero, S.One}
|
||
|
|
assert set(dens.subs({N: 3, m: 2, n: 1}).doit().values()) == {Rational(1, 3), Rational(2, 3)}
|
||
|
|
k = Dummy('k', integer=True)
|
||
|
|
assert expec.dummy_eq(
|
||
|
|
Sum(Piecewise((k*binomial(m, k)*binomial(N - m, -k + n)
|
||
|
|
/binomial(N, n), k > 2), (0, True)), (k, 0, n)))
|
||
|
|
|
||
|
|
def test_rademacher():
|
||
|
|
X = Rademacher('X')
|
||
|
|
t = Symbol('t')
|
||
|
|
|
||
|
|
assert E(X) == 0
|
||
|
|
assert variance(X) == 1
|
||
|
|
assert density(X)[-1] == S.Half
|
||
|
|
assert density(X)[1] == S.Half
|
||
|
|
assert characteristic_function(X)(t) == exp(I*t)/2 + exp(-I*t)/2
|
||
|
|
assert moment_generating_function(X)(t) == exp(t) / 2 + exp(-t) / 2
|
||
|
|
|
||
|
|
def test_ideal_soliton():
|
||
|
|
raises(ValueError, lambda : IdealSoliton('sol', -12))
|
||
|
|
raises(ValueError, lambda : IdealSoliton('sol', 13.2))
|
||
|
|
raises(ValueError, lambda : IdealSoliton('sol', 0))
|
||
|
|
f = Function('f')
|
||
|
|
raises(ValueError, lambda : density(IdealSoliton('sol', 10)).pmf(f))
|
||
|
|
|
||
|
|
k = Symbol('k', integer=True, positive=True)
|
||
|
|
x = Symbol('x', integer=True, positive=True)
|
||
|
|
t = Symbol('t')
|
||
|
|
sol = IdealSoliton('sol', k)
|
||
|
|
assert density(sol).low == S.One
|
||
|
|
assert density(sol).high == k
|
||
|
|
assert density(sol).dict == Density(density(sol))
|
||
|
|
assert density(sol).pmf(x) == Piecewise((1/k, Eq(x, 1)), (1/(x*(x - 1)), k >= x), (0, True))
|
||
|
|
|
||
|
|
k_vals = [5, 20, 50, 100, 1000]
|
||
|
|
for i in k_vals:
|
||
|
|
assert E(sol.subs(k, i)) == harmonic(i) == moment(sol.subs(k, i), 1)
|
||
|
|
assert variance(sol.subs(k, i)) == (i - 1) + harmonic(i) - harmonic(i)**2 == cmoment(sol.subs(k, i),2)
|
||
|
|
assert skewness(sol.subs(k, i)) == smoment(sol.subs(k, i), 3)
|
||
|
|
assert kurtosis(sol.subs(k, i)) == smoment(sol.subs(k, i), 4)
|
||
|
|
|
||
|
|
assert exp(I*t)/10 + Sum(exp(I*t*x)/(x*x - x), (x, 2, k)).subs(k, 10).doit() == characteristic_function(sol.subs(k, 10))(t)
|
||
|
|
assert exp(t)/10 + Sum(exp(t*x)/(x*x - x), (x, 2, k)).subs(k, 10).doit() == moment_generating_function(sol.subs(k, 10))(t)
|
||
|
|
|
||
|
|
def test_robust_soliton():
|
||
|
|
raises(ValueError, lambda : RobustSoliton('robSol', -12, 0.1, 0.02))
|
||
|
|
raises(ValueError, lambda : RobustSoliton('robSol', 13, 1.89, 0.1))
|
||
|
|
raises(ValueError, lambda : RobustSoliton('robSol', 15, 0.6, -2.31))
|
||
|
|
f = Function('f')
|
||
|
|
raises(ValueError, lambda : density(RobustSoliton('robSol', 15, 0.6, 0.1)).pmf(f))
|
||
|
|
|
||
|
|
k = Symbol('k', integer=True, positive=True)
|
||
|
|
delta = Symbol('delta', positive=True)
|
||
|
|
c = Symbol('c', positive=True)
|
||
|
|
robSol = RobustSoliton('robSol', k, delta, c)
|
||
|
|
assert density(robSol).low == 1
|
||
|
|
assert density(robSol).high == k
|
||
|
|
|
||
|
|
k_vals = [10, 20, 50]
|
||
|
|
delta_vals = [0.2, 0.4, 0.6]
|
||
|
|
c_vals = [0.01, 0.03, 0.05]
|
||
|
|
for x in k_vals:
|
||
|
|
for y in delta_vals:
|
||
|
|
for z in c_vals:
|
||
|
|
assert E(robSol.subs({k: x, delta: y, c: z})) == moment(robSol.subs({k: x, delta: y, c: z}), 1)
|
||
|
|
assert variance(robSol.subs({k: x, delta: y, c: z})) == cmoment(robSol.subs({k: x, delta: y, c: z}), 2)
|
||
|
|
assert skewness(robSol.subs({k: x, delta: y, c: z})) == smoment(robSol.subs({k: x, delta: y, c: z}), 3)
|
||
|
|
assert kurtosis(robSol.subs({k: x, delta: y, c: z})) == smoment(robSol.subs({k: x, delta: y, c: z}), 4)
|
||
|
|
|
||
|
|
def test_FiniteRV():
|
||
|
|
F = FiniteRV('F', {1: S.Half, 2: Rational(1, 4), 3: Rational(1, 4)}, check=True)
|
||
|
|
p = Symbol("p", positive=True)
|
||
|
|
|
||
|
|
assert dict(density(F).items()) == {S.One: S.Half, S(2): Rational(1, 4), S(3): Rational(1, 4)}
|
||
|
|
assert P(F >= 2) == S.Half
|
||
|
|
assert quantile(F)(p) == Piecewise((nan, p > S.One), (S.One, p <= S.Half),\
|
||
|
|
(S(2), p <= Rational(3, 4)),(S(3), True))
|
||
|
|
|
||
|
|
assert pspace(F).domain.as_boolean() == Or(
|
||
|
|
*[Eq(F.symbol, i) for i in [1, 2, 3]])
|
||
|
|
|
||
|
|
assert F.pspace.domain.set == FiniteSet(1, 2, 3)
|
||
|
|
raises(ValueError, lambda: FiniteRV('F', {1: S.Half, 2: S.Half, 3: S.Half}, check=True))
|
||
|
|
raises(ValueError, lambda: FiniteRV('F', {1: S.Half, 2: Rational(-1, 2), 3: S.One}, check=True))
|
||
|
|
raises(ValueError, lambda: FiniteRV('F', {1: S.One, 2: Rational(3, 2), 3: S.Zero,\
|
||
|
|
4: Rational(-1, 2), 5: Rational(-3, 4), 6: Rational(-1, 4)}, check=True))
|
||
|
|
|
||
|
|
# purposeful invalid pmf but it should not raise since check=False
|
||
|
|
# see test_drv_types.test_ContinuousRV for explanation
|
||
|
|
X = FiniteRV('X', {1: 1, 2: 2})
|
||
|
|
assert E(X) == 5
|
||
|
|
assert P(X <= 2) + P(X > 2) != 1
|
||
|
|
|
||
|
|
def test_density_call():
|
||
|
|
from sympy.abc import p
|
||
|
|
x = Bernoulli('x', p)
|
||
|
|
d = density(x)
|
||
|
|
assert d(0) == 1 - p
|
||
|
|
assert d(S.Zero) == 1 - p
|
||
|
|
assert d(5) == 0
|
||
|
|
|
||
|
|
assert 0 in d
|
||
|
|
assert 5 not in d
|
||
|
|
assert d(S.Zero) == d[S.Zero]
|
||
|
|
|
||
|
|
|
||
|
|
def test_DieDistribution():
|
||
|
|
from sympy.abc import x
|
||
|
|
X = DieDistribution(6)
|
||
|
|
assert X.pmf(S.Half) is S.Zero
|
||
|
|
assert X.pmf(x).subs({x: 1}).doit() == Rational(1, 6)
|
||
|
|
assert X.pmf(x).subs({x: 7}).doit() == 0
|
||
|
|
assert X.pmf(x).subs({x: -1}).doit() == 0
|
||
|
|
assert X.pmf(x).subs({x: Rational(1, 3)}).doit() == 0
|
||
|
|
raises(ValueError, lambda: X.pmf(Matrix([0, 0])))
|
||
|
|
raises(ValueError, lambda: X.pmf(x**2 - 1))
|
||
|
|
|
||
|
|
def test_FinitePSpace():
|
||
|
|
X = Die('X', 6)
|
||
|
|
space = pspace(X)
|
||
|
|
assert space.density == DieDistribution(6)
|
||
|
|
|
||
|
|
def test_symbolic_conditions():
|
||
|
|
B = Bernoulli('B', Rational(1, 4))
|
||
|
|
D = Die('D', 4)
|
||
|
|
b, n = symbols('b, n')
|
||
|
|
Y = P(Eq(B, b))
|
||
|
|
Z = E(D > n)
|
||
|
|
assert Y == \
|
||
|
|
Piecewise((Rational(1, 4), Eq(b, 1)), (0, True)) + \
|
||
|
|
Piecewise((Rational(3, 4), Eq(b, 0)), (0, True))
|
||
|
|
assert Z == \
|
||
|
|
Piecewise((Rational(1, 4), n < 1), (0, True)) + Piecewise((S.Half, n < 2), (0, True)) + \
|
||
|
|
Piecewise((Rational(3, 4), n < 3), (0, True)) + Piecewise((S.One, n < 4), (0, True))
|