# Source code for sympy.matrices.expressions.blockmatrix

from matexpr import MatrixExpr, ZeroMatrix, Identity
from matmul import MatMul
from transpose import Transpose
from trace import Trace
from inverse import Inverse
from sympy.matrices import Matrix, eye
from sympy import Tuple, Basic, sympify, FiniteSet, Add

[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]

"""
is_BlockMatrix = True
is_BlockDiagMatrix = False
def __new__(cls, mat):
if not isinstance(mat, Matrix):
mat = Matrix(mat)
data = Tuple(*mat.mat)
shape = Tuple(*sympify(mat.shape))
obj = Basic.__new__(cls, data, shape)
obj.mat = mat
return obj

@property
def shape(self):
numrows = numcols = 0
M = self.mat
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.mat.shape

@property
def blocks(self):
return self.mat

@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 _blockmul(self, other):

if  (other.is_Matrix and other.is_BlockMatrix and
self.blockshape[1] == other.blockshape[0] and
self.colblocksizes == other.rowblocksizes):
return BlockMatrix(self.mat*other.mat)

return MatrixExpr.__mul__(self, other)

if  (other.is_Matrix and other.is_BlockMatrix and
self.blockshape == other.blockshape and
self.rowblocksizes == other.rowblocksizes and
self.colblocksizes == other.colblocksizes):
return BlockMatrix(self.mat + other.mat)

def _eval_transpose(self):
# Flip all the individual matrices
matrices = [Transpose(matrix) for matrix in self.mat.mat]
# Make a copy
mat = Matrix(self.blockshape[0], self.blockshape[1], matrices)
# Transpose the block structure
mat = mat.transpose()
return BlockMatrix(mat)

def _eval_trace(self):
if self.rowblocksizes == self.colblocksizes:
for i in range(self.blockshape[0])])
raise NotImplementedError("Can't perform trace of irregular blockshape")

#def transpose(self):
#    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, (Inverse(self.blocks[0]), ))
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*Inverse(A22)*A21
C2 = A22 - A21*Inverse(A11)*A12
mat = Matrix([[Inverse(C1), Inverse(-A11)*A12*Inverse(C2)],
[-Inverse(C2)*A21*Inverse(A11), Inverse(C2)]])
return BlockMatrix(mat)
else:
raise NotImplementedError()

def inverse(self):
return Inverse(self)

def _entry(self, i, j):
idx = 0
# 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.mat[i,j].is_Identity:
return False
if i!=j and not self.mat[i,j].is_ZeroMatrix:
return False
return True
@property
def is_structurally_symmetric(self):
return self.rowblocksizes == self.colblocksizes

[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]

"""
is_BlockDiagMatrix = True
def __new__(cls, *mats):
data_matrix = eye(len(mats))
for i, mat in enumerate(mats):
data_matrix[i,i] = mat

for r in range(len(mats)):
for c in range(len(mats)):
if r == c:
continue
n = mats[r].rows
m = mats[c].cols
data_matrix[r, c] = ZeroMatrix(n, m)

shape = Tuple(*sympify(mat.shape))
data = Tuple(*data_matrix.mat)
obj = Basic.__new__(cls, data, shape, Tuple(*mats))
obj.mat = data_matrix
return obj

@property
def diag(self):
return self.args[2]

def _eval_inverse(self):
return BlockDiagMatrix(*[Inverse(mat) for mat in self.diag])

def _blockmul(self, other):
if  (other.is_Matrix and other.is_BlockMatrix and
other.is_BlockDiagMatrix and
self.blockshape[1] == other.blockshape[0] and
self.colblocksizes == other.rowblocksizes):
return BlockDiagMatrix(*[a*b for a, b in zip(self.diag,other.diag)])
else:
return BlockMatrix._blockmul(self, other)

if  (other.is_Matrix and other.is_BlockMatrix and
other.is_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.diag,other.diag)])
else:

[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]
"""
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)

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:
block = blocks[0]
for b in blocks[1:]:
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)

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