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