Source code for sympy.matrices.expressions.blockmatrix

from __future__ import print_function, division

from sympy import ask, Q
from sympy.core import Basic, Add, sympify
from sympy.core.compatibility import range
from sympy.strategies import typed, exhaust, condition, do_one, unpack
from sympy.strategies.traverse import bottom_up
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.transpose import Transpose, transpose
from sympy.matrices.expressions.trace import Trace
from sympy.matrices.expressions.determinant import det, Determinant
from sympy.matrices.expressions.slice import MatrixSlice
from sympy.matrices.expressions.inverse import Inverse
from sympy.matrices import Matrix, ShapeError
from sympy.functions.elementary.complexes import re, im

[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) Matrix([ [X, Z], [0, Y]]) >>> C = BlockMatrix([[Identity(n), Z]]) >>> print(C) Matrix([[I, Z]]) >>> print(block_collapse(C*B)) Matrix([[X, Z*Y + Z]]) """ def __new__(cls, *args): from sympy.matrices.immutable import ImmutableDenseMatrix args = map(sympify, args) mat = ImmutableDenseMatrix(*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") def _eval_determinant(self): if self.blockshape == (2, 2): [[A, B], [C, D]] = self.blocks.tolist() if ask(Q.invertible(A)): return det(A)*det(D - C*A.I*B) elif ask(Q.invertible(D)): return det(D)*det(A - B*D.I*C) return Determinant(self) def as_real_imag(self): real_matrices = [re(matrix) for matrix in self.blocks] real_matrices = Matrix(self.blockshape[0], self.blockshape[1], real_matrices) im_matrices = [im(matrix) for matrix in self.blocks] im_matrices = Matrix(self.blockshape[0], self.blockshape[1], im_matrices) return (real_matrices, im_matrices)
[docs] def transpose(self): """Return transpose of matrix. Examples ======== >>> from sympy import MatrixSymbol, BlockMatrix, ZeroMatrix >>> from 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() Matrix([ [X.T, 0], [Z.T, Y.T]]) >>> _.transpose() Matrix([ [X, Z], [0, Y]]) """ return self._eval_transpose()
def _entry(self, i, j): # Find row entry for row_block, numrows in enumerate(self.rowblocksizes): if (i < numrows) != False: break else: i -= numrows for col_block, numcols in enumerate(self.colblocksizes): if (j < numcols) != False: 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) Matrix([ [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 ImmutableDenseMatrix 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 ImmutableDenseMatrix(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) Matrix([ [X, Z], [0, Y]]) >>> C = BlockMatrix([[Identity(n), Z]]) >>> print(C) Matrix([[I, Z]]) >>> print(block_collapse(C*B)) Matrix([[X, Z*Y + Z]]) """ hasbm = lambda expr: isinstance(expr, MatrixExpr) and expr.has(BlockMatrix) rule = exhaust( bottom_up(exhaust(condition(hasbm, typed( {MatAdd: do_one(bc_matadd, bc_block_plus_ident), MatMul: do_one(bc_matmul, bc_dist), Transpose: bc_transpose, Inverse: bc_inverse, BlockMatrix: do_one(bc_unpack, deblock)}))))) 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) elif isinstance(A, BlockMatrix): matrices[i] = A._blockmul(BlockMatrix([[B]])) matrices.pop(i+1) elif isinstance(B, BlockMatrix): matrices[i] = BlockMatrix([[A]])._blockmul(B) matrices.pop(i+1) else: i+=1 return MatMul(factor, *matrices).doit() def bc_transpose(expr): return BlockMatrix(block_collapse(expr.arg).blocks.applyfunc(transpose).T) def bc_inverse(expr): expr2 = blockinverse_1x1(expr) if expr != expr2: return expr2 return blockinverse_2x2(Inverse(reblock_2x2(expr.arg))) def blockinverse_1x1(expr): if isinstance(expr.arg, BlockMatrix) and expr.arg.blockshape == (1, 1): mat = Matrix([[expr.arg.blocks[0].inverse()]]) return BlockMatrix(mat) return expr def blockinverse_2x2(expr): if isinstance(expr.arg, BlockMatrix) and expr.arg.blockshape == (2, 2): # Cite: The Matrix Cookbook Section 9.1.3 [[A, B], [C, D]] = expr.arg.blocks.tolist() return BlockMatrix([[ (A - B*D.I*C).I, (-A).I*B*(D - C*A.I*B).I], [-(D - C*A.I*B).I*C*A.I, (D - C*A.I*B).I]]) else: return expr def deblock(B): """ Flatten a BlockMatrix of BlockMatrices """ if not isinstance(B, BlockMatrix) or not B.blocks.has(BlockMatrix): return B wrap = lambda x: x if isinstance(x, BlockMatrix) else BlockMatrix([[x]]) bb = B.blocks.applyfunc(wrap) # everything is a block from sympy import Matrix try: MM = Matrix(0, sum(bb[0, i].blocks.shape[1] for i in range(bb.shape[1])), []) for row in range(0, bb.shape[0]): M = Matrix(bb[row, 0].blocks) for col in range(1, bb.shape[1]): M = M.row_join(bb[row, col].blocks) MM = MM.col_join(M) return BlockMatrix(MM) except ShapeError: return B def reblock_2x2(B): """ Reblock a BlockMatrix so that it has 2x2 blocks of block matrices """ if not isinstance(B, BlockMatrix) or not all(d > 2 for d in B.blocks.shape): return B BM = BlockMatrix # for brevity's sake return BM([[ B.blocks[0, 0], BM(B.blocks[0, 1:])], [BM(B.blocks[1:, 0]), BM(B.blocks[1:, 1:])]]) def bounds(sizes): """ Convert sequence of numbers into pairs of low-high pairs >>> from sympy.matrices.expressions.blockmatrix import bounds >>> bounds((1, 10, 50)) [(0, 1), (1, 11), (11, 61)] """ low = 0 rv = [] for size in sizes: rv.append((low, low + size)) low += size return rv def blockcut(expr, rowsizes, colsizes): """ Cut a matrix expression into Blocks >>> from sympy import ImmutableMatrix, blockcut >>> M = ImmutableMatrix(4, 4, range(16)) >>> B = blockcut(M, (1, 3), (1, 3)) >>> type(B).__name__ 'BlockMatrix' >>> ImmutableMatrix(B.blocks[0, 1]) Matrix([[1, 2, 3]]) """ rowbounds = bounds(rowsizes) colbounds = bounds(colsizes) return BlockMatrix([[MatrixSlice(expr, rowbound, colbound) for colbound in colbounds] for rowbound in rowbounds])