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):
rv = []
nonrv = []
for a in expr.args:
if is_random(a):
rv.append(a)
else:
nonrv.append(a)
return Mul.fromiter(nonrv) * Expectation(Mul.fromiter(rv),
condition=condition)
return self
def expand(self, **hints):
expr = self.args[0]
condition = self._condition
if not is_random(expr):
return expr
if isinstance(expr, Add):
return Add(*[
Expectation(a, condition=condition).expand() for a in expr.args
])
elif isinstance(expr, (Mul, MatMul)):
if isinstance(_expand(expr), Add):
return Expectation(_expand(expr)).expand()
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
def expand(self, **hints):
expr = self.args[0]
condition = self._condition
if not is_random(expr):
return expr
if isinstance(expr, Add):
return Add(*[
Expectation(a, condition=condition).expand() for a in expr.args
])
elif isinstance(expr, Mul):
if isinstance(_expand(expr), Add):
return Expectation(_expand(expr)).expand()
rv = []
nonrv = []
for a in expr.args:
if is_random(a):
rv.append(a)
else:
nonrv.append(a)
return Mul(*nonrv) * Expectation(Mul(*rv), condition=condition)
return self
def expand(expr):
from sympy import expand as _expand
if isinstance(expr, Tuple):
return expr
expr = _expand(expr)
_, args = expr.as_coeff_add()
args = list(args)
for i in range(len(args)):
c,m = args[i].as_coeff_mul()
for o in m:
c = c*o
args[i] = c
return Add(*args)
def expand(expr):
"""
Expand an expression by making sure that Mul objects are evaluated by using
the * operator (which might be overloaded by the types in the expression).
"""
from sympy import expand as _expand
if isinstance(expr, Tuple):
return expr
coeff, args = _expand(expr, MatMul=True).as_coeff_add()
newargs = [
c * reduce(mul, m, 1) for a in args for c, m in [a.as_coeff_mul()]
]
expr = coeff
for e in newargs:
expr += e
return expr
