import itertools
from sympy import Expr, Add, Mul, S, Integral, Eq, Sum, Symbol
from sympy.core.compatibility import default_sort_key
from sympy.core.evaluate import global_evaluate
from sympy.core.sympify import _sympify
from sympy.stats import variance, covariance
from sympy.stats.rv import RandomSymbol, probability, expectation
__all__ = ['Probability', 'Expectation', 'Variance', 'Covariance']
[docs]class Probability(Expr):
"""
Symbolic expression for the probability.
Examples
========
>>> from sympy.stats import Probability, Normal
>>> from sympy import Integral
>>> X = Normal("X", 0, 1)
>>> prob = Probability(X > 1)
>>> prob
Probability(X > 1)
Integral representation:
>>> prob.rewrite(Integral)
Integral(sqrt(2)*exp(-_z**2/2)/(2*sqrt(pi)), (_z, 1, oo))
Evaluation of the integral:
>>> prob.evaluate_integral()
sqrt(2)*(-sqrt(2)*sqrt(pi)*erf(sqrt(2)/2) + sqrt(2)*sqrt(pi))/(4*sqrt(pi))
"""
def __new__(cls, prob, condition=None, **kwargs):
prob = _sympify(prob)
if condition is None:
obj = Expr.__new__(cls, prob)
else:
condition = _sympify(condition)
obj = Expr.__new__(cls, prob, condition)
obj._condition = condition
return obj
def _eval_rewrite_as_Integral(self, arg, condition=None, **kwargs):
return probability(arg, condition, evaluate=False)
def _eval_rewrite_as_Sum(self, arg, condition=None, **kwargs):
return probability(arg, condition, evaluate=False)
def evaluate_integral(self):
return self.rewrite(Integral).doit()
[docs]class Expectation(Expr):
"""
Symbolic expression for the expectation.
Examples
========
>>> from sympy.stats import Expectation, Normal, Probability
>>> from sympy import symbols, Integral
>>> mu = symbols("mu")
>>> sigma = symbols("sigma", positive=True)
>>> X = Normal("X", mu, sigma)
>>> Expectation(X)
Expectation(X)
>>> Expectation(X).evaluate_integral().simplify()
mu
To get the integral expression of the expectation:
>>> Expectation(X).rewrite(Integral)
Integral(sqrt(2)*X*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo))
The same integral expression, in more abstract terms:
>>> Expectation(X).rewrite(Probability)
Integral(x*Probability(Eq(X, x)), (x, -oo, oo))
This class is aware of some properties of the expectation:
>>> from sympy.abc import a
>>> Expectation(a*X)
Expectation(a*X)
>>> Y = Normal("Y", 0, 1)
>>> Expectation(X + Y)
Expectation(X + Y)
To expand the ``Expectation`` into its expression, use ``doit()``:
>>> Expectation(X + Y).doit()
Expectation(X) + Expectation(Y)
>>> Expectation(a*X + Y).doit()
a*Expectation(X) + Expectation(Y)
>>> Expectation(a*X + Y)
Expectation(a*X + Y)
"""
def __new__(cls, expr, condition=None, **kwargs):
expr = _sympify(expr)
if condition is None:
if not expr.has(RandomSymbol):
return expr
obj = Expr.__new__(cls, expr)
else:
condition = _sympify(condition)
obj = Expr.__new__(cls, expr, condition)
obj._condition = condition
return obj
def doit(self, **hints):
expr = self.args[0]
condition = self._condition
if not expr.has(RandomSymbol):
return expr
if isinstance(expr, Add):
return Add(*[Expectation(a, condition=condition).doit() for a in expr.args])
elif isinstance(expr, Mul):
rv = []
nonrv = []
for a in expr.args:
if isinstance(a, RandomSymbol) or a.has(RandomSymbol):
rv.append(a)
else:
nonrv.append(a)
return Mul(*nonrv)*Expectation(Mul(*rv), condition=condition)
return self
def _eval_rewrite_as_Probability(self, arg, condition=None, **kwargs):
rvs = arg.atoms(RandomSymbol)
if len(rvs) > 1:
raise NotImplementedError()
if len(rvs) == 0:
return arg
rv = rvs.pop()
if rv.pspace is None:
raise ValueError("Probability space not known")
symbol = rv.symbol
if symbol.name[0].isupper():
symbol = Symbol(symbol.name.lower())
else :
symbol = Symbol(symbol.name + "_1")
if rv.pspace.is_Continuous:
return Integral(arg.replace(rv, symbol)*Probability(Eq(rv, symbol), condition), (symbol, rv.pspace.domain.set.inf, rv.pspace.domain.set.sup))
else:
if rv.pspace.is_Finite:
raise NotImplementedError
else:
return Sum(arg.replace(rv, symbol)*Probability(Eq(rv, symbol), condition), (symbol, rv.pspace.domain.set.inf, rv.pspace.set.sup))
def _eval_rewrite_as_Integral(self, arg, condition=None, **kwargs):
return expectation(arg, condition=condition, evaluate=False)
def _eval_rewrite_as_Sum(self, arg, condition=None, **kwargs):
return self.rewrite(Integral)
def evaluate_integral(self):
return self.rewrite(Integral).doit()
[docs]class Variance(Expr):
"""
Symbolic expression for the variance.
Examples
========
>>> from sympy import symbols, Integral
>>> from sympy.stats import Normal, Expectation, Variance, Probability
>>> mu = symbols("mu", positive=True)
>>> sigma = symbols("sigma", positive=True)
>>> X = Normal("X", mu, sigma)
>>> Variance(X)
Variance(X)
>>> Variance(X).evaluate_integral()
sigma**2
Integral representation of the underlying calculations:
>>> Variance(X).rewrite(Integral)
Integral(sqrt(2)*(X - Integral(sqrt(2)*X*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo)))**2*exp(-(X - mu)**2/(2*sigma**2))/(2*sqrt(pi)*sigma), (X, -oo, oo))
Integral representation, without expanding the PDF:
>>> Variance(X).rewrite(Probability)
-Integral(x*Probability(Eq(X, x)), (x, -oo, oo))**2 + Integral(x**2*Probability(Eq(X, x)), (x, -oo, oo))
Rewrite the variance in terms of the expectation
>>> Variance(X).rewrite(Expectation)
-Expectation(X)**2 + Expectation(X**2)
Some transformations based on the properties of the variance may happen:
>>> from sympy.abc import a
>>> Y = Normal("Y", 0, 1)
>>> Variance(a*X)
Variance(a*X)
To expand the variance in its expression, use ``doit()``:
>>> Variance(a*X).doit()
a**2*Variance(X)
>>> Variance(X + Y)
Variance(X + Y)
>>> Variance(X + Y).doit()
2*Covariance(X, Y) + Variance(X) + Variance(Y)
"""
def __new__(cls, arg, condition=None, **kwargs):
arg = _sympify(arg)
if condition is None:
obj = Expr.__new__(cls, arg)
else:
condition = _sympify(condition)
obj = Expr.__new__(cls, arg, condition)
obj._condition = condition
return obj
def doit(self, **hints):
arg = self.args[0]
condition = self._condition
if not arg.has(RandomSymbol):
return S.Zero
if isinstance(arg, RandomSymbol):
return self
elif isinstance(arg, Add):
rv = []
for a in arg.args:
if a.has(RandomSymbol):
rv.append(a)
variances = Add(*map(lambda xv: Variance(xv, condition).doit(), rv))
map_to_covar = lambda x: 2*Covariance(*x, condition=condition).doit()
covariances = Add(*map(map_to_covar, itertools.combinations(rv, 2)))
return variances + covariances
elif isinstance(arg, Mul):
nonrv = []
rv = []
for a in arg.args:
if a.has(RandomSymbol):
rv.append(a)
else:
nonrv.append(a**2)
if len(rv) == 0:
return S.Zero
return Mul(*nonrv)*Variance(Mul(*rv), condition)
# this expression contains a RandomSymbol somehow:
return self
def _eval_rewrite_as_Expectation(self, arg, condition=None, **kwargs):
e1 = Expectation(arg**2, condition)
e2 = Expectation(arg, condition)**2
return e1 - e2
def _eval_rewrite_as_Probability(self, arg, condition=None, **kwargs):
return self.rewrite(Expectation).rewrite(Probability)
def _eval_rewrite_as_Integral(self, arg, condition=None, **kwargs):
return variance(self.args[0], self._condition, evaluate=False)
def _eval_rewrite_as_Sum(self, arg, condition=None, **kwargs):
return self.rewrite(Integral)
def evaluate_integral(self):
return self.rewrite(Integral).doit()
[docs]class Covariance(Expr):
"""
Symbolic expression for the covariance.
Examples
========
>>> from sympy.stats import Covariance
>>> from sympy.stats import Normal
>>> X = Normal("X", 3, 2)
>>> Y = Normal("Y", 0, 1)
>>> Z = Normal("Z", 0, 1)
>>> W = Normal("W", 0, 1)
>>> cexpr = Covariance(X, Y)
>>> cexpr
Covariance(X, Y)
Evaluate the covariance, `X` and `Y` are independent,
therefore zero is the result:
>>> cexpr.evaluate_integral()
0
Rewrite the covariance expression in terms of expectations:
>>> from sympy.stats import Expectation
>>> cexpr.rewrite(Expectation)
Expectation(X*Y) - Expectation(X)*Expectation(Y)
In order to expand the argument, use ``doit()``:
>>> from sympy.abc import a, b, c, d
>>> Covariance(a*X + b*Y, c*Z + d*W)
Covariance(a*X + b*Y, c*Z + d*W)
>>> Covariance(a*X + b*Y, c*Z + d*W).doit()
a*c*Covariance(X, Z) + a*d*Covariance(W, X) + b*c*Covariance(Y, Z) + b*d*Covariance(W, Y)
This class is aware of some properties of the covariance:
>>> Covariance(X, X).doit()
Variance(X)
>>> Covariance(a*X, b*Y).doit()
a*b*Covariance(X, Y)
"""
def __new__(cls, arg1, arg2, condition=None, **kwargs):
arg1 = _sympify(arg1)
arg2 = _sympify(arg2)
if kwargs.pop('evaluate', global_evaluate[0]):
arg1, arg2 = sorted([arg1, arg2], key=default_sort_key)
if condition is None:
obj = Expr.__new__(cls, arg1, arg2)
else:
condition = _sympify(condition)
obj = Expr.__new__(cls, arg1, arg2, condition)
obj._condition = condition
return obj
def doit(self, **hints):
arg1 = self.args[0]
arg2 = self.args[1]
condition = self._condition
if arg1 == arg2:
return Variance(arg1, condition).doit()
if not arg1.has(RandomSymbol):
return S.Zero
if not arg2.has(RandomSymbol):
return S.Zero
arg1, arg2 = sorted([arg1, arg2], key=default_sort_key)
if isinstance(arg1, RandomSymbol) and isinstance(arg2, RandomSymbol):
return Covariance(arg1, arg2, condition)
coeff_rv_list1 = self._expand_single_argument(arg1.expand())
coeff_rv_list2 = self._expand_single_argument(arg2.expand())
addends = [a*b*Covariance(*sorted([r1, r2], key=default_sort_key), condition=condition)
for (a, r1) in coeff_rv_list1 for (b, r2) in coeff_rv_list2]
return Add(*addends)
@classmethod
def _expand_single_argument(cls, expr):
# return (coefficient, random_symbol) pairs:
if isinstance(expr, RandomSymbol):
return [(S.One, expr)]
elif isinstance(expr, Add):
outval = []
for a in expr.args:
if isinstance(a, Mul):
outval.append(cls._get_mul_nonrv_rv_tuple(a))
elif isinstance(a, RandomSymbol):
outval.append((S.One, a))
return outval
elif isinstance(expr, Mul):
return [cls._get_mul_nonrv_rv_tuple(expr)]
elif expr.has(RandomSymbol):
return [(S.One, expr)]
@classmethod
def _get_mul_nonrv_rv_tuple(cls, m):
rv = []
nonrv = []
for a in m.args:
if a.has(RandomSymbol):
rv.append(a)
else:
nonrv.append(a)
return (Mul(*nonrv), Mul(*rv))
def _eval_rewrite_as_Expectation(self, arg1, arg2, condition=None, **kwargs):
e1 = Expectation(arg1*arg2, condition)
e2 = Expectation(arg1, condition)*Expectation(arg2, condition)
return e1 - e2
def _eval_rewrite_as_Probability(self, arg1, arg2, condition=None, **kwargs):
return self.rewrite(Expectation).rewrite(Probability)
def _eval_rewrite_as_Integral(self, arg1, arg2, condition=None, **kwargs):
return covariance(self.args[0], self.args[1], self._condition, evaluate=False)
def _eval_rewrite_as_Sum(self, arg1, arg2, condition=None, **kwargs):
return self.rewrite(Integral)
def evaluate_integral(self):
return self.rewrite(Integral).doit()
