# Source code for sympy.matrices.expressions.blockmatrix

from __future__ import print_function, division

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.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

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:
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()
return det(A)*det(D - C*A.I*B)
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 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()
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)

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:

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

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:]:
if nonblocks:
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 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])