Source code for sympy.matrices.expressions.blockmatrix

from sympy.core import Tuple, Basic, Add
from sympy.rules import typed, canon, debug, do_one, unpack
from sympy.functions import transpose
from sympy.utilities import sift

from sympy.matrices.expressions.matexpr import MatrixExpr, ZeroMatrix, Identity
from sympy.matrices.expressions.matmul import MatMul
from sympy.matrices.expressions.matadd import MatAdd
from sympy.matrices.expressions.matpow import MatPow
from sympy.matrices.expressions.transpose import Transpose
from sympy.matrices.expressions.trace import Trace
from sympy.matrices.expressions.inverse import Inverse
from sympy.matrices import Matrix, eye


[docs]class BlockMatrix(MatrixExpr): """A BlockMatrix is a Matrix composed of other smaller, submatrices The submatrices are stored in a SymPy Matrix object but accessed as part of a Matrix Expression >>> from sympy import (MatrixSymbol, BlockMatrix, symbols, ... Identity, 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] """ def __new__(cls, *args): from sympy.matrices.immutable import ImmutableMatrix mat = ImmutableMatrix(*args) obj = Basic.__new__(cls, mat) return obj @property def shape(self): numrows = numcols = 0 M = self.blocks for i in range(M.shape[0]): numrows += M[i, 0].shape[0] for i in range(M.shape[1]): numcols += M[0, i].shape[1] return (numrows, numcols) @property def blockshape(self): return self.blocks.shape @property def blocks(self): return self.args[0] @property def rowblocksizes(self): return [self.blocks[i, 0].rows for i in range(self.blockshape[0])] @property def colblocksizes(self): return [self.blocks[0, i].cols for i in range(self.blockshape[1])] def structurally_equal(self, other): return (isinstance(other, BlockMatrix) and self.shape == other.shape and self.blockshape == other.blockshape and self.rowblocksizes == other.rowblocksizes and self.colblocksizes == other.colblocksizes) def _blockmul(self, other): if (isinstance(other, BlockMatrix) and self.colblocksizes == other.rowblocksizes): return BlockMatrix(self.blocks*other.blocks) return self * other def _blockadd(self, other): if (isinstance(other, BlockMatrix) and self.structurally_equal(other)): return BlockMatrix(self.blocks + other.blocks) return self + other def _eval_transpose(self): # Flip all the individual matrices matrices = [transpose(matrix) for matrix in self.blocks] # Make a copy M = Matrix(self.blockshape[0], self.blockshape[1], matrices) # Transpose the block structure M = M.transpose() return BlockMatrix(M) def _eval_trace(self): if self.rowblocksizes == self.colblocksizes: return Add(*[Trace(self.blocks[i, i]) for i in range(self.blockshape[0])]) raise NotImplementedError( "Can't perform trace of irregular blockshape")
[docs] def transpose(self): """Return transpose of matrix. Examples ======== >>> from sympy import MatrixSymbol, BlockMatrix, ZeroMatrix >>> from sympy.abc import l, m, n >>> X = MatrixSymbol('X', n, n) >>> Y = MatrixSymbol('Y', m ,m) >>> Z = MatrixSymbol('Z', n, m) >>> B = BlockMatrix([[X, Z], [ZeroMatrix(m,n), Y]]) >>> B.transpose() [X', 0] [Z', Y'] >>> _.transpose() [X, Z] [0, Y] """ return self._eval_transpose()
def _eval_inverse(self, expand=False): # Inverse of one by one block matrix is easy if self.blockshape == (1, 1): mat = Matrix(1, 1, (self.blocks[0].inverse(),)) return BlockMatrix(mat) # Inverse of a two by two block matrix is known elif expand and self.blockshape == (2, 2): # Cite: The Matrix Cookbook Section 9.1.3 A11, A12, A21, A22 = (self.blocks[0, 0], self.blocks[0, 1], self.blocks[1, 0], self.blocks[1, 1]) C1 = A11 - A12*A22.I*A21 C2 = A22 - A21*A11.I*A12 mat = Matrix([[C1.I, (-A11).I*A12*C2.I], [-C2.I*A21*A11.I, C2.I]]) return BlockMatrix(mat) else: return Inverse(self)
[docs] def inverse(self, expand=False): """Return inverse of matrix. Examples ======== >>> from sympy import MatrixSymbol, BlockMatrix, ZeroMatrix >>> from sympy.abc import l, m, n >>> X = MatrixSymbol('X', n, n) >>> BlockMatrix([[X]]).inverse() [X^-1] >>> Y = MatrixSymbol('Y', m ,m) >>> Z = MatrixSymbol('Z', n, m) >>> B = BlockMatrix([[X, Z], [ZeroMatrix(m,n), Y]]) >>> B [X, Z] [0, Y] >>> B.inverse(expand=True) [X^-1, (-1)*X^-1*Z*Y^-1] [ 0, Y^-1] """ return self._eval_inverse(expand)
def _entry(self, i, j): # Find row entry for row_block, numrows in enumerate(self.rowblocksizes): if i < numrows: break else: i -= numrows for col_block, numcols in enumerate(self.colblocksizes): if j < numcols: break else: j -= numcols return self.blocks[row_block, col_block][i, j] @property def is_Identity(self): if self.blockshape[0] != self.blockshape[1]: return False for i in range(self.blockshape[0]): for j in range(self.blockshape[1]): if i==j and not self.blocks[i, j].is_Identity: return False if i!=j and not self.blocks[i, j].is_ZeroMatrix: return False return True @property def is_structurally_symmetric(self): return self.rowblocksizes == self.colblocksizes def equals(self, other): if self == other: return True if (isinstance(other, BlockMatrix) and self.blocks == other.blocks): return True return super(BlockMatrix, self).equals(other)
[docs]class BlockDiagMatrix(BlockMatrix): """ A BlockDiagMatrix is a BlockMatrix with matrices only along the diagonal >>> from sympy import MatrixSymbol, BlockDiagMatrix, symbols, Identity >>> n,m,l = symbols('n m l') >>> X = MatrixSymbol('X', n, n) >>> Y = MatrixSymbol('Y', m ,m) >>> BlockDiagMatrix(X, Y) [X, 0] [0, Y] """ def __new__(cls, *mats): return Basic.__new__(BlockDiagMatrix, *mats) @property def diag(self): return self.args @property def blocks(self): from sympy.matrices.immutable import ImmutableMatrix mats = self.args data = [[mats[i] if i == j else ZeroMatrix(mats[i].rows, mats[j].cols) for j in range(len(mats))] for i in range(len(mats))] return ImmutableMatrix(data) @property def shape(self): return (sum(block.rows for block in self.args), sum(block.cols for block in self.args)) @property def blockshape(self): n = len(self.args) return (n, n) @property def rowblocksizes(self): return [block.rows for block in self.args] @property def colblocksizes(self): return [block.cols for block in self.args] def _eval_inverse(self, expand='ignored'): return BlockDiagMatrix(*[mat.inverse() for mat in self.args]) def _blockmul(self, other): if (isinstance(other, BlockDiagMatrix) and self.colblocksizes == other.rowblocksizes): return BlockDiagMatrix(*[a*b for a, b in zip(self.args, other.args)]) else: return BlockMatrix._blockmul(self, other) def _blockadd(self, other): if (isinstance(other, BlockDiagMatrix) and self.blockshape == other.blockshape and self.rowblocksizes == other.rowblocksizes and self.colblocksizes == other.colblocksizes): return BlockDiagMatrix(*[a + b for a, b in zip(self.args, other.args)]) else: return BlockMatrix._blockadd(self, other)
[docs]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] """ rule = canon(typed({MatAdd: do_one(bc_matadd, bc_block_plus_ident), MatMul: do_one(bc_matmul, bc_dist), BlockMatrix: bc_unpack})) result = rule(expr) try: return result.doit() except AttributeError: return result
def bc_unpack(expr): if expr.blockshape == (1, 1): return expr.blocks[0, 0] return expr def bc_matadd(expr): args = sift(expr.args, lambda M: isinstance(M, BlockMatrix)) blocks = args[True] if not blocks: return expr nonblocks = args[False] block = blocks[0] for b in blocks[1:]: block = block._blockadd(b) if nonblocks: return MatAdd(*nonblocks) + block else: return 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 isinstance(arg, 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).doit() return expr def bc_dist(expr): """ Turn a*[X, Y] into [a*X, a*Y] """ factor, mat = expr.as_coeff_mmul() if factor != 1 and isinstance(unpack(mat), BlockMatrix): B = unpack(mat).blocks return BlockMatrix([[factor * B[i, j] for j in range(B.cols)] for i in range(B.rows)]) return expr def bc_matmul(expr): factor, matrices = expr.as_coeff_matrices() i = 0 while (i+1 < len(matrices)): A, B = matrices[i:i+2] if isinstance(A, BlockMatrix) and isinstance(B, BlockMatrix): matrices[i] = A._blockmul(B) matrices.pop(i+1) else: i+=1 return MatMul(factor, *matrices).doit()