def bc_matadd(expr):
nonblocks = [arg for arg in expr.args if not arg.is_BlockMatrix]
blocks = [arg for arg in expr.args if arg.is_BlockMatrix]
if not blocks:
return expr
block = blocks[0]
for b in blocks[1:]:
block = block._blockadd(b)
return MatAdd(*(nonblocks + [block]))
def bc_block_plus_ident(expr):
idents = [arg for arg in expr.args if arg.is_Identity]
if not idents:
return expr
blocks = [arg for arg in expr.args if arg.is_BlockMatrix]
if (blocks and all(b.structurally_equal(blocks[0]) for b in blocks)
and blocks[0].is_structurally_symmetric):
block_id = BlockDiagMatrix(
*[Identity(k) for k in blocks[0].rowblocksizes])
return MatAdd(block_id * len(idents), *blocks)
return expr
def transpose(self):
if isinstance(self, Transpose):
return self.arg
if self.is_Mul:
return MatMul(*[Transpose(arg) for arg in self.args[::-1]])
if self.is_Add:
return MatAdd(*[Transpose(arg) for arg in self.args])
try:
return self._eval_transpose()
except (AttributeError, NotImplementedError):
return Basic.__new__(Transpose, self)
def __new__(cls, mat):
if not mat.is_Matrix:
return mat
if isinstance(mat, Transpose):
return mat.arg
if hasattr(mat, 'transpose'):
return mat.transpose()
if mat.is_Mul:
return MatMul(*[Transpose(arg) for arg in mat.args[::-1]])
if mat.is_Add:
return MatAdd(*[Transpose(arg) for arg in mat.args])
return Basic.__new__(cls, mat)
def __new__(cls, *args):
# Check that the shape of the args is consistent
matrices = [arg for arg in args if arg.is_Matrix]
for i in range(len(matrices) - 1):
A, B = matrices[i:i + 2]
if A.cols != B.rows:
raise ShapeError("Matrices %s and %s are not aligned" % (A, B))
if any(arg.is_zero for arg in args):
return ZeroMatrix(matrices[0].rows, matrices[-1].cols)
expr = matrixify(Mul.__new__(cls, *args))
if expr.is_Add:
return MatAdd(*expr.args)
if expr.is_Pow:
assert expr.exp.is_Integer
expr = Basic.__new__(MatMul, *[expr.base for i in range(expr.exp)])
if not expr.is_Mul:
return expr
if any(arg.is_Matrix and arg.is_ZeroMatrix for arg in expr.args):
return ZeroMatrix(*expr.shape)
# Clear out Identities
nonmats = [M for M in expr.args if not M.is_Matrix] # scalars
mats = [M for M in expr.args if M.is_Matrix] # matrices
if any(M.is_Identity for M in mats): # Any identities around?
newmats = [M for M in mats if not M.is_Identity] # clear out
if len(newmats) == 0: # Did we lose everything?
newmats = [Identity(expr.rows)] # put just one back in
if mats != newmats: # Removed some I's but not everything?
return MatMul(*(nonmats + newmats)) # Repeat with simpler expr
return expr
def __rsub__(self, other):
return MatAdd(other, -self).doit()
def __sub__(self, other):
return MatAdd(self, -other).doit()
def __radd__(self, other):
return MatAdd(other, self).doit()
def __add__(self, other):
return MatAdd(self, other).doit()
def block_collapse(expr):
"""Evaluates a block matrix expression
>>> from sympy import MatrixSymbol, BlockMatrix, symbols, Identity, Matrix, ZeroMatrix, block_collapse
>>> n,m,l = symbols('n m l')
>>> X = MatrixSymbol('X', n, n)
>>> Y = MatrixSymbol('Y', m ,m)
>>> Z = MatrixSymbol('Z', n, m)
>>> B = BlockMatrix([[X, Z], [ZeroMatrix(m, n), Y]])
>>> print B
[X, Z]
[0, Y]
>>> C = BlockMatrix([[Identity(n), Z]])
>>> print C
[I, Z]
>>> print block_collapse(C*B)
[X, Z + Z*Y]
"""
if expr.__class__ in [tuple, list, set, frozenset]:
return expr.__class__([block_collapse(arg) for arg in expr])
if expr.__class__ in [Tuple, FiniteSet]:
return expr.__class__(*[block_collapse(arg) for arg in expr])
if not expr.is_Matrix or (not expr.is_Add and not expr.is_Mul
and not expr.is_Transpose and not expr.is_Pow
and not expr.is_Inverse):
return expr
if expr.is_Transpose:
expr = Transpose(block_collapse(expr.arg))
if expr.is_Transpose and expr.arg.is_BlockMatrix:
expr = expr.arg.eval_transpose()
return expr
if expr.is_Inverse:
return Inverse(block_collapse(expr.arg))
# Recurse on the subargs
args = list(expr.args)
for i in range(len(args)):
arg = args[i]
newarg = block_collapse(arg)
while(newarg != arg): # Repeat until no new changes
arg = newarg
newarg = block_collapse(arg)
args[i] = newarg
if tuple(args) != expr.args:
expr = expr.__class__(*args)
# Turn -[X, Y] into [-X, -Y]
if (expr.is_Mul and len(expr.args)==2 and not expr.args[0].is_Matrix
and expr.args[1].is_BlockMatrix):
if expr.args[1].is_BlockDiagMatrix:
return BlockDiagMatrix(
*[expr.args[0]*arg for arg in expr.args[1].diag])
else:
return BlockMatrix(expr.args[0]*expr.args[1].mat)
if expr.is_Add:
nonblocks = [arg for arg in expr.args if not arg.is_BlockMatrix]
blocks = [arg for arg in expr.args if arg.is_BlockMatrix]
if not blocks:
return MatAdd(*nonblocks)
block = blocks[0]
for b in blocks[1:]:
block = block._blockadd(b)
if block.blockshape == (1,1):
# Bring all the non-blocks into the block_matrix
mat = Matrix(1, 1, (block.blocks[0,0] + MatAdd(*nonblocks), ))
return BlockMatrix(mat)
# Add identities to the blocks as block identities
for i, mat in enumerate(nonblocks):
c, M = mat.as_coeff_Mul()
if M.is_Identity and block.is_structurally_symmetric:
block_id = BlockDiagMatrix(
*[c*Identity(k) for k in block.rowblocksizes])
nonblocks.pop(i)
block = block._blockadd(block_id)
return MatAdd(*(nonblocks+[block]))
if expr.is_Mul:
nonmatrices = [arg for arg in expr.args if not arg.is_Matrix]
matrices = [arg for arg in expr.args if arg.is_Matrix]
i = 0
while (i+1 < len(matrices)):
A, B = matrices[i:i+2]
if A.is_BlockMatrix and B.is_BlockMatrix:
matrices[i] = A._blockmul(B)
matrices.pop(i+1)
else:
i+=1
return MatMul(*(nonmatrices + matrices))
if expr.is_Pow:
rv = expr.base
for i in range(1, expr.exp):
rv = rv._blockmul(expr.base)
return rv
