309 lines
10 KiB
Python
309 lines
10 KiB
Python
import itertools
|
|
|
|
from sympy.core.add import Add
|
|
from sympy.core.expr import Expr
|
|
from sympy.core.function import expand as _expand
|
|
from sympy.core.mul import Mul
|
|
from sympy.core.singleton import S
|
|
from sympy.matrices.common import ShapeError
|
|
from sympy.matrices.expressions.matexpr import MatrixExpr
|
|
from sympy.matrices.expressions.matmul import MatMul
|
|
from sympy.matrices.expressions.special import ZeroMatrix
|
|
from sympy.stats.rv import RandomSymbol, is_random
|
|
from sympy.core.sympify import _sympify
|
|
from sympy.stats.symbolic_probability import Variance, Covariance, Expectation
|
|
|
|
|
|
class ExpectationMatrix(Expectation, MatrixExpr):
|
|
"""
|
|
Expectation of a random matrix expression.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.stats import ExpectationMatrix, Normal
|
|
>>> from sympy.stats.rv import RandomMatrixSymbol
|
|
>>> from sympy import symbols, MatrixSymbol, Matrix
|
|
>>> k = symbols("k")
|
|
>>> A, B = MatrixSymbol("A", k, k), MatrixSymbol("B", k, k)
|
|
>>> X, Y = RandomMatrixSymbol("X", k, 1), RandomMatrixSymbol("Y", k, 1)
|
|
>>> ExpectationMatrix(X)
|
|
ExpectationMatrix(X)
|
|
>>> ExpectationMatrix(A*X).shape
|
|
(k, 1)
|
|
|
|
To expand the expectation in its expression, use ``expand()``:
|
|
|
|
>>> ExpectationMatrix(A*X + B*Y).expand()
|
|
A*ExpectationMatrix(X) + B*ExpectationMatrix(Y)
|
|
>>> ExpectationMatrix((X + Y)*(X - Y).T).expand()
|
|
ExpectationMatrix(X*X.T) - ExpectationMatrix(X*Y.T) + ExpectationMatrix(Y*X.T) - ExpectationMatrix(Y*Y.T)
|
|
|
|
To evaluate the ``ExpectationMatrix``, use ``doit()``:
|
|
|
|
>>> N11, N12 = Normal('N11', 11, 1), Normal('N12', 12, 1)
|
|
>>> N21, N22 = Normal('N21', 21, 1), Normal('N22', 22, 1)
|
|
>>> M11, M12 = Normal('M11', 1, 1), Normal('M12', 2, 1)
|
|
>>> M21, M22 = Normal('M21', 3, 1), Normal('M22', 4, 1)
|
|
>>> x1 = Matrix([[N11, N12], [N21, N22]])
|
|
>>> x2 = Matrix([[M11, M12], [M21, M22]])
|
|
>>> ExpectationMatrix(x1 + x2).doit()
|
|
Matrix([
|
|
[12, 14],
|
|
[24, 26]])
|
|
|
|
"""
|
|
def __new__(cls, expr, condition=None):
|
|
expr = _sympify(expr)
|
|
if condition is None:
|
|
if not is_random(expr):
|
|
return expr
|
|
obj = Expr.__new__(cls, expr)
|
|
else:
|
|
condition = _sympify(condition)
|
|
obj = Expr.__new__(cls, expr, condition)
|
|
|
|
obj._shape = expr.shape
|
|
obj._condition = condition
|
|
return obj
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._shape
|
|
|
|
def expand(self, **hints):
|
|
expr = self.args[0]
|
|
condition = self._condition
|
|
if not is_random(expr):
|
|
return expr
|
|
|
|
if isinstance(expr, Add):
|
|
return Add.fromiter(Expectation(a, condition=condition).expand()
|
|
for a in expr.args)
|
|
|
|
expand_expr = _expand(expr)
|
|
if isinstance(expand_expr, Add):
|
|
return Add.fromiter(Expectation(a, condition=condition).expand()
|
|
for a in expand_expr.args)
|
|
|
|
elif isinstance(expr, (Mul, MatMul)):
|
|
rv = []
|
|
nonrv = []
|
|
postnon = []
|
|
|
|
for a in expr.args:
|
|
if is_random(a):
|
|
if rv:
|
|
rv.extend(postnon)
|
|
else:
|
|
nonrv.extend(postnon)
|
|
postnon = []
|
|
rv.append(a)
|
|
elif a.is_Matrix:
|
|
postnon.append(a)
|
|
else:
|
|
nonrv.append(a)
|
|
|
|
# In order to avoid infinite-looping (MatMul may call .doit() again),
|
|
# do not rebuild
|
|
if len(nonrv) == 0:
|
|
return self
|
|
return Mul.fromiter(nonrv)*Expectation(Mul.fromiter(rv),
|
|
condition=condition)*Mul.fromiter(postnon)
|
|
|
|
return self
|
|
|
|
class VarianceMatrix(Variance, MatrixExpr):
|
|
"""
|
|
Variance of a random matrix probability expression. Also known as
|
|
Covariance matrix, auto-covariance matrix, dispersion matrix,
|
|
or variance-covariance matrix.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.stats import VarianceMatrix
|
|
>>> from sympy.stats.rv import RandomMatrixSymbol
|
|
>>> from sympy import symbols, MatrixSymbol
|
|
>>> k = symbols("k")
|
|
>>> A, B = MatrixSymbol("A", k, k), MatrixSymbol("B", k, k)
|
|
>>> X, Y = RandomMatrixSymbol("X", k, 1), RandomMatrixSymbol("Y", k, 1)
|
|
>>> VarianceMatrix(X)
|
|
VarianceMatrix(X)
|
|
>>> VarianceMatrix(X).shape
|
|
(k, k)
|
|
|
|
To expand the variance in its expression, use ``expand()``:
|
|
|
|
>>> VarianceMatrix(A*X).expand()
|
|
A*VarianceMatrix(X)*A.T
|
|
>>> VarianceMatrix(A*X + B*Y).expand()
|
|
2*A*CrossCovarianceMatrix(X, Y)*B.T + A*VarianceMatrix(X)*A.T + B*VarianceMatrix(Y)*B.T
|
|
"""
|
|
def __new__(cls, arg, condition=None):
|
|
arg = _sympify(arg)
|
|
|
|
if 1 not in arg.shape:
|
|
raise ShapeError("Expression is not a vector")
|
|
|
|
shape = (arg.shape[0], arg.shape[0]) if arg.shape[1] == 1 else (arg.shape[1], arg.shape[1])
|
|
|
|
if condition:
|
|
obj = Expr.__new__(cls, arg, condition)
|
|
else:
|
|
obj = Expr.__new__(cls, arg)
|
|
|
|
obj._shape = shape
|
|
obj._condition = condition
|
|
return obj
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._shape
|
|
|
|
def expand(self, **hints):
|
|
arg = self.args[0]
|
|
condition = self._condition
|
|
|
|
if not is_random(arg):
|
|
return ZeroMatrix(*self.shape)
|
|
|
|
if isinstance(arg, RandomSymbol):
|
|
return self
|
|
elif isinstance(arg, Add):
|
|
rv = []
|
|
for a in arg.args:
|
|
if is_random(a):
|
|
rv.append(a)
|
|
variances = Add(*(Variance(xv, condition).expand() for xv in rv))
|
|
map_to_covar = lambda x: 2*Covariance(*x, condition=condition).expand()
|
|
covariances = Add(*map(map_to_covar, itertools.combinations(rv, 2)))
|
|
return variances + covariances
|
|
elif isinstance(arg, (Mul, MatMul)):
|
|
nonrv = []
|
|
rv = []
|
|
for a in arg.args:
|
|
if is_random(a):
|
|
rv.append(a)
|
|
else:
|
|
nonrv.append(a)
|
|
if len(rv) == 0:
|
|
return ZeroMatrix(*self.shape)
|
|
# Avoid possible infinite loops with MatMul:
|
|
if len(nonrv) == 0:
|
|
return self
|
|
# Variance of many multiple matrix products is not implemented:
|
|
if len(rv) > 1:
|
|
return self
|
|
return Mul.fromiter(nonrv)*Variance(Mul.fromiter(rv),
|
|
condition)*(Mul.fromiter(nonrv)).transpose()
|
|
|
|
# this expression contains a RandomSymbol somehow:
|
|
return self
|
|
|
|
class CrossCovarianceMatrix(Covariance, MatrixExpr):
|
|
"""
|
|
Covariance of a random matrix probability expression.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.stats import CrossCovarianceMatrix
|
|
>>> from sympy.stats.rv import RandomMatrixSymbol
|
|
>>> from sympy import symbols, MatrixSymbol
|
|
>>> k = symbols("k")
|
|
>>> A, B = MatrixSymbol("A", k, k), MatrixSymbol("B", k, k)
|
|
>>> C, D = MatrixSymbol("C", k, k), MatrixSymbol("D", k, k)
|
|
>>> X, Y = RandomMatrixSymbol("X", k, 1), RandomMatrixSymbol("Y", k, 1)
|
|
>>> Z, W = RandomMatrixSymbol("Z", k, 1), RandomMatrixSymbol("W", k, 1)
|
|
>>> CrossCovarianceMatrix(X, Y)
|
|
CrossCovarianceMatrix(X, Y)
|
|
>>> CrossCovarianceMatrix(X, Y).shape
|
|
(k, k)
|
|
|
|
To expand the covariance in its expression, use ``expand()``:
|
|
|
|
>>> CrossCovarianceMatrix(X + Y, Z).expand()
|
|
CrossCovarianceMatrix(X, Z) + CrossCovarianceMatrix(Y, Z)
|
|
>>> CrossCovarianceMatrix(A*X, Y).expand()
|
|
A*CrossCovarianceMatrix(X, Y)
|
|
>>> CrossCovarianceMatrix(A*X, B.T*Y).expand()
|
|
A*CrossCovarianceMatrix(X, Y)*B
|
|
>>> CrossCovarianceMatrix(A*X + B*Y, C.T*Z + D.T*W).expand()
|
|
A*CrossCovarianceMatrix(X, W)*D + A*CrossCovarianceMatrix(X, Z)*C + B*CrossCovarianceMatrix(Y, W)*D + B*CrossCovarianceMatrix(Y, Z)*C
|
|
|
|
"""
|
|
def __new__(cls, arg1, arg2, condition=None):
|
|
arg1 = _sympify(arg1)
|
|
arg2 = _sympify(arg2)
|
|
|
|
if (1 not in arg1.shape) or (1 not in arg2.shape) or (arg1.shape[1] != arg2.shape[1]):
|
|
raise ShapeError("Expression is not a vector")
|
|
|
|
shape = (arg1.shape[0], arg2.shape[0]) if arg1.shape[1] == 1 and arg2.shape[1] == 1 \
|
|
else (1, 1)
|
|
|
|
if condition:
|
|
obj = Expr.__new__(cls, arg1, arg2, condition)
|
|
else:
|
|
obj = Expr.__new__(cls, arg1, arg2)
|
|
|
|
obj._shape = shape
|
|
obj._condition = condition
|
|
return obj
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._shape
|
|
|
|
def expand(self, **hints):
|
|
arg1 = self.args[0]
|
|
arg2 = self.args[1]
|
|
condition = self._condition
|
|
|
|
if arg1 == arg2:
|
|
return VarianceMatrix(arg1, condition).expand()
|
|
|
|
if not is_random(arg1) or not is_random(arg2):
|
|
return ZeroMatrix(*self.shape)
|
|
|
|
if isinstance(arg1, RandomSymbol) and isinstance(arg2, RandomSymbol):
|
|
return CrossCovarianceMatrix(arg1, arg2, condition)
|
|
|
|
coeff_rv_list1 = self._expand_single_argument(arg1.expand())
|
|
coeff_rv_list2 = self._expand_single_argument(arg2.expand())
|
|
|
|
addends = [a*CrossCovarianceMatrix(r1, r2, condition=condition)*b.transpose()
|
|
for (a, r1) in coeff_rv_list1 for (b, r2) in coeff_rv_list2]
|
|
return Add.fromiter(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, MatMul)):
|
|
outval.append(cls._get_mul_nonrv_rv_tuple(a))
|
|
elif is_random(a):
|
|
outval.append((S.One, a))
|
|
|
|
return outval
|
|
elif isinstance(expr, (Mul, MatMul)):
|
|
return [cls._get_mul_nonrv_rv_tuple(expr)]
|
|
elif is_random(expr):
|
|
return [(S.One, expr)]
|
|
|
|
@classmethod
|
|
def _get_mul_nonrv_rv_tuple(cls, m):
|
|
rv = []
|
|
nonrv = []
|
|
for a in m.args:
|
|
if is_random(a):
|
|
rv.append(a)
|
|
else:
|
|
nonrv.append(a)
|
|
return (Mul.fromiter(nonrv), Mul.fromiter(rv))
|