123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969 |
- from sympy.assumptions.ask import (Q, ask)
- from sympy.core import Basic, Add, Mul, S
- from sympy.core.sympify import _sympify
- from sympy.functions.elementary.complexes import re, im
- from sympy.strategies import typed, exhaust, condition, do_one, unpack
- from sympy.strategies.traverse import bottom_up
- from sympy.utilities.iterables import is_sequence, sift
- from sympy.utilities.misc import filldedent
- from sympy.matrices import Matrix, ShapeError
- from sympy.matrices.common import NonInvertibleMatrixError
- from sympy.matrices.expressions.determinant import det, Determinant
- from sympy.matrices.expressions.inverse import Inverse
- from sympy.matrices.expressions.matadd import MatAdd
- from sympy.matrices.expressions.matexpr import MatrixExpr, MatrixElement
- from sympy.matrices.expressions.matmul import MatMul
- from sympy.matrices.expressions.matpow import MatPow
- from sympy.matrices.expressions.slice import MatrixSlice
- from sympy.matrices.expressions.special import ZeroMatrix, Identity
- from sympy.matrices.expressions.trace import trace
- from sympy.matrices.expressions.transpose import Transpose, transpose
- class BlockMatrix(MatrixExpr):
- """A BlockMatrix is a Matrix comprised of other matrices.
- 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 + Z*Y]])
- Some matrices might be comprised of rows of blocks with
- the matrices in each row having the same height and the
- rows all having the same total number of columns but
- not having the same number of columns for each matrix
- in each row. In this case, the matrix is not a block
- matrix and should be instantiated by Matrix.
- >>> from sympy import ones, Matrix
- >>> dat = [
- ... [ones(3,2), ones(3,3)*2],
- ... [ones(2,3)*3, ones(2,2)*4]]
- ...
- >>> BlockMatrix(dat)
- Traceback (most recent call last):
- ...
- ValueError:
- Although this matrix is comprised of blocks, the blocks do not fill
- the matrix in a size-symmetric fashion. To create a full matrix from
- these arguments, pass them directly to Matrix.
- >>> Matrix(dat)
- Matrix([
- [1, 1, 2, 2, 2],
- [1, 1, 2, 2, 2],
- [1, 1, 2, 2, 2],
- [3, 3, 3, 4, 4],
- [3, 3, 3, 4, 4]])
- See Also
- ========
- sympy.matrices.matrices.MatrixBase.irregular
- """
- def __new__(cls, *args, **kwargs):
- from sympy.matrices.immutable import ImmutableDenseMatrix
- isMat = lambda i: getattr(i, 'is_Matrix', False)
- if len(args) != 1 or \
- not is_sequence(args[0]) or \
- len({isMat(r) for r in args[0]}) != 1:
- raise ValueError(filldedent('''
- expecting a sequence of 1 or more rows
- containing Matrices.'''))
- rows = args[0] if args else []
- if not isMat(rows):
- if rows and isMat(rows[0]):
- rows = [rows] # rows is not list of lists or []
- # regularity check
- # same number of matrices in each row
- blocky = ok = len({len(r) for r in rows}) == 1
- if ok:
- # same number of rows for each matrix in a row
- for r in rows:
- ok = len({i.rows for i in r}) == 1
- if not ok:
- break
- blocky = ok
- if ok:
- # same number of cols for each matrix in each col
- for c in range(len(rows[0])):
- ok = len({rows[i][c].cols
- for i in range(len(rows))}) == 1
- if not ok:
- break
- if not ok:
- # same total cols in each row
- ok = len({
- sum([i.cols for i in r]) for r in rows}) == 1
- if blocky and ok:
- raise ValueError(filldedent('''
- Although this matrix is comprised of blocks,
- the blocks do not fill the matrix in a
- size-symmetric fashion. To create a full matrix
- from these arguments, pass them directly to
- Matrix.'''))
- raise ValueError(filldedent('''
- When there are not the same number of rows in each
- row's matrices or there are not the same number of
- total columns in each row, the matrix is not a
- block matrix. If this matrix is known to consist of
- blocks fully filling a 2-D space then see
- Matrix.irregular.'''))
- mat = ImmutableDenseMatrix(rows, evaluate=False)
- 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 == (1, 1):
- return det(self.blocks[0, 0])
- 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)
- def transpose(self):
- """Return transpose of matrix.
- Examples
- ========
- >>> from sympy import MatrixSymbol, BlockMatrix, ZeroMatrix
- >>> from sympy.abc import 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 schur(self, mat = 'A', generalized = False):
- """Return the Schur Complement of the 2x2 BlockMatrix
- Parameters
- ==========
- mat : String, optional
- The matrix with respect to which the
- Schur Complement is calculated. 'A' is
- used by default
- generalized : bool, optional
- If True, returns the generalized Schur
- Component which uses Moore-Penrose Inverse
- Examples
- ========
- >>> from sympy import symbols, MatrixSymbol, BlockMatrix
- >>> m, n = symbols('m n')
- >>> A = MatrixSymbol('A', n, n)
- >>> B = MatrixSymbol('B', n, m)
- >>> C = MatrixSymbol('C', m, n)
- >>> D = MatrixSymbol('D', m, m)
- >>> X = BlockMatrix([[A, B], [C, D]])
- The default Schur Complement is evaluated with "A"
- >>> X.schur()
- -C*A**(-1)*B + D
- >>> X.schur('D')
- A - B*D**(-1)*C
- Schur complement with non-invertible matrices is not
- defined. Instead, the generalized Schur complement can
- be calculated which uses the Moore-Penrose Inverse. To
- achieve this, `generalized` must be set to `True`
- >>> X.schur('B', generalized=True)
- C - D*(B.T*B)**(-1)*B.T*A
- >>> X.schur('C', generalized=True)
- -A*(C.T*C)**(-1)*C.T*D + B
- Returns
- =======
- M : Matrix
- The Schur Complement Matrix
- Raises
- ======
- ShapeError
- If the block matrix is not a 2x2 matrix
- NonInvertibleMatrixError
- If given matrix is non-invertible
- References
- ==========
- .. [1] Wikipedia Article on Schur Component : https://en.wikipedia.org/wiki/Schur_complement
- See Also
- ========
- sympy.matrices.matrices.MatrixBase.pinv
- """
- if self.blockshape == (2, 2):
- [[A, B],
- [C, D]] = self.blocks.tolist()
- d={'A' : A, 'B' : B, 'C' : C, 'D' : D}
- try:
- inv = (d[mat].T*d[mat]).inv()*d[mat].T if generalized else d[mat].inv()
- if mat == 'A':
- return D - C * inv * B
- elif mat == 'B':
- return C - D * inv * A
- elif mat == 'C':
- return B - A * inv * D
- elif mat == 'D':
- return A - B * inv * C
- #For matrices where no sub-matrix is square
- return self
- except NonInvertibleMatrixError:
- raise NonInvertibleMatrixError('The given matrix is not invertible. Please set generalized=True \
- to compute the generalized Schur Complement which uses Moore-Penrose Inverse')
- else:
- raise ShapeError('Schur Complement can only be calculated for 2x2 block matrices')
- def LDUdecomposition(self):
- """Returns the Block LDU decomposition of
- a 2x2 Block Matrix
- Returns
- =======
- (L, D, U) : Matrices
- L : Lower Diagonal Matrix
- D : Diagonal Matrix
- U : Upper Diagonal Matrix
- Examples
- ========
- >>> from sympy import symbols, MatrixSymbol, BlockMatrix, block_collapse
- >>> m, n = symbols('m n')
- >>> A = MatrixSymbol('A', n, n)
- >>> B = MatrixSymbol('B', n, m)
- >>> C = MatrixSymbol('C', m, n)
- >>> D = MatrixSymbol('D', m, m)
- >>> X = BlockMatrix([[A, B], [C, D]])
- >>> L, D, U = X.LDUdecomposition()
- >>> block_collapse(L*D*U)
- Matrix([
- [A, B],
- [C, D]])
- Raises
- ======
- ShapeError
- If the block matrix is not a 2x2 matrix
- NonInvertibleMatrixError
- If the matrix "A" is non-invertible
- See Also
- ========
- sympy.matrices.expressions.blockmatrix.BlockMatrix.UDLdecomposition
- sympy.matrices.expressions.blockmatrix.BlockMatrix.LUdecomposition
- """
- if self.blockshape == (2,2):
- [[A, B],
- [C, D]] = self.blocks.tolist()
- try:
- AI = A.I
- except NonInvertibleMatrixError:
- raise NonInvertibleMatrixError('Block LDU decomposition cannot be calculated when\
- "A" is singular')
- Ip = Identity(B.shape[0])
- Iq = Identity(B.shape[1])
- Z = ZeroMatrix(*B.shape)
- L = BlockMatrix([[Ip, Z], [C*AI, Iq]])
- D = BlockDiagMatrix(A, self.schur())
- U = BlockMatrix([[Ip, AI*B],[Z.T, Iq]])
- return L, D, U
- else:
- raise ShapeError("Block LDU decomposition is supported only for 2x2 block matrices")
- def UDLdecomposition(self):
- """Returns the Block UDL decomposition of
- a 2x2 Block Matrix
- Returns
- =======
- (U, D, L) : Matrices
- U : Upper Diagonal Matrix
- D : Diagonal Matrix
- L : Lower Diagonal Matrix
- Examples
- ========
- >>> from sympy import symbols, MatrixSymbol, BlockMatrix, block_collapse
- >>> m, n = symbols('m n')
- >>> A = MatrixSymbol('A', n, n)
- >>> B = MatrixSymbol('B', n, m)
- >>> C = MatrixSymbol('C', m, n)
- >>> D = MatrixSymbol('D', m, m)
- >>> X = BlockMatrix([[A, B], [C, D]])
- >>> U, D, L = X.UDLdecomposition()
- >>> block_collapse(U*D*L)
- Matrix([
- [A, B],
- [C, D]])
- Raises
- ======
- ShapeError
- If the block matrix is not a 2x2 matrix
- NonInvertibleMatrixError
- If the matrix "D" is non-invertible
- See Also
- ========
- sympy.matrices.expressions.blockmatrix.BlockMatrix.LDUdecomposition
- sympy.matrices.expressions.blockmatrix.BlockMatrix.LUdecomposition
- """
- if self.blockshape == (2,2):
- [[A, B],
- [C, D]] = self.blocks.tolist()
- try:
- DI = D.I
- except NonInvertibleMatrixError:
- raise NonInvertibleMatrixError('Block UDL decomposition cannot be calculated when\
- "D" is singular')
- Ip = Identity(A.shape[0])
- Iq = Identity(B.shape[1])
- Z = ZeroMatrix(*B.shape)
- U = BlockMatrix([[Ip, B*DI], [Z.T, Iq]])
- D = BlockDiagMatrix(self.schur('D'), D)
- L = BlockMatrix([[Ip, Z],[DI*C, Iq]])
- return U, D, L
- else:
- raise ShapeError("Block UDL decomposition is supported only for 2x2 block matrices")
- def LUdecomposition(self):
- """Returns the Block LU decomposition of
- a 2x2 Block Matrix
- Returns
- =======
- (L, U) : Matrices
- L : Lower Diagonal Matrix
- U : Upper Diagonal Matrix
- Examples
- ========
- >>> from sympy import symbols, MatrixSymbol, BlockMatrix, block_collapse
- >>> m, n = symbols('m n')
- >>> A = MatrixSymbol('A', n, n)
- >>> B = MatrixSymbol('B', n, m)
- >>> C = MatrixSymbol('C', m, n)
- >>> D = MatrixSymbol('D', m, m)
- >>> X = BlockMatrix([[A, B], [C, D]])
- >>> L, U = X.LUdecomposition()
- >>> block_collapse(L*U)
- Matrix([
- [A, B],
- [C, D]])
- Raises
- ======
- ShapeError
- If the block matrix is not a 2x2 matrix
- NonInvertibleMatrixError
- If the matrix "A" is non-invertible
- See Also
- ========
- sympy.matrices.expressions.blockmatrix.BlockMatrix.UDLdecomposition
- sympy.matrices.expressions.blockmatrix.BlockMatrix.LDUdecomposition
- """
- if self.blockshape == (2,2):
- [[A, B],
- [C, D]] = self.blocks.tolist()
- try:
- A = A**0.5
- AI = A.I
- except NonInvertibleMatrixError:
- raise NonInvertibleMatrixError('Block LU decomposition cannot be calculated when\
- "A" is singular')
- Z = ZeroMatrix(*B.shape)
- Q = self.schur()**0.5
- L = BlockMatrix([[A, Z], [C*AI, Q]])
- U = BlockMatrix([[A, AI*B],[Z.T, Q]])
- return L, U
- else:
- raise ShapeError("Block LU decomposition is supported only for 2x2 block matrices")
- def _entry(self, i, j, **kwargs):
- # Find row entry
- orig_i, orig_j = i, j
- for row_block, numrows in enumerate(self.rowblocksizes):
- cmp = i < numrows
- if cmp == True:
- break
- elif cmp == False:
- i -= numrows
- elif row_block < self.blockshape[0] - 1:
- # Can't tell which block and it's not the last one, return unevaluated
- return MatrixElement(self, orig_i, orig_j)
- for col_block, numcols in enumerate(self.colblocksizes):
- cmp = j < numcols
- if cmp == True:
- break
- elif cmp == False:
- j -= numcols
- elif col_block < self.blockshape[1] - 1:
- return MatrixElement(self, orig_i, orig_j)
- 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().equals(other)
- class BlockDiagMatrix(BlockMatrix):
- """A sparse matrix with block matrices along its diagonals
- Examples
- ========
- >>> from sympy import MatrixSymbol, BlockDiagMatrix, symbols
- >>> 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]])
- Notes
- =====
- If you want to get the individual diagonal blocks, use
- :meth:`get_diag_blocks`.
- See Also
- ========
- sympy.matrices.dense.diag
- """
- def __new__(cls, *mats):
- return Basic.__new__(BlockDiagMatrix, *[_sympify(m) for m in 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, evaluate=False)
- @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 _all_square_blocks(self):
- """Returns true if all blocks are square"""
- return all(mat.is_square for mat in self.args)
- def _eval_determinant(self):
- if self._all_square_blocks():
- return Mul(*[det(mat) for mat in self.args])
- # At least one block is non-square. Since the entire matrix must be square we know there must
- # be at least two blocks in this matrix, in which case the entire matrix is necessarily rank-deficient
- return S.Zero
- def _eval_inverse(self, expand='ignored'):
- if self._all_square_blocks():
- return BlockDiagMatrix(*[mat.inverse() for mat in self.args])
- # See comment in _eval_determinant()
- raise NonInvertibleMatrixError('Matrix det == 0; not invertible.')
- def _eval_transpose(self):
- return BlockDiagMatrix(*[mat.transpose() 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)
- def get_diag_blocks(self):
- """Return the list of diagonal blocks of the matrix.
- Examples
- ========
- >>> from sympy import BlockDiagMatrix, Matrix
- >>> A = Matrix([[1, 2], [3, 4]])
- >>> B = Matrix([[5, 6], [7, 8]])
- >>> M = BlockDiagMatrix(A, B)
- How to get diagonal blocks from the block diagonal matrix:
- >>> diag_blocks = M.get_diag_blocks()
- >>> diag_blocks[0]
- Matrix([
- [1, 2],
- [3, 4]])
- >>> diag_blocks[1]
- Matrix([
- [5, 6],
- [7, 8]])
- """
- return self.args
- def block_collapse(expr):
- """Evaluates a block 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 + Z*Y]])
- """
- from sympy.strategies.util import expr_fns
- hasbm = lambda expr: isinstance(expr, MatrixExpr) and expr.has(BlockMatrix)
- conditioned_rl = condition(
- hasbm,
- typed(
- {MatAdd: do_one(bc_matadd, bc_block_plus_ident),
- MatMul: do_one(bc_matmul, bc_dist),
- MatPow: bc_matmul,
- Transpose: bc_transpose,
- Inverse: bc_inverse,
- BlockMatrix: do_one(bc_unpack, deblock)}
- )
- )
- rule = exhaust(
- bottom_up(
- exhaust(conditioned_rl),
- fns=expr_fns
- )
- )
- result = rule(expr)
- doit = getattr(result, 'doit', None)
- if doit is not None:
- return doit()
- else:
- 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])
- rest = [arg for arg in expr.args if not arg.is_Identity and not isinstance(arg, BlockMatrix)]
- return MatAdd(block_id * len(idents), *blocks, *rest).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:
- return expr
- unpacked = unpack(mat)
- if isinstance(unpacked, BlockDiagMatrix):
- B = unpacked.diag
- new_B = [factor * mat for mat in B]
- return BlockDiagMatrix(*new_B)
- elif isinstance(unpacked, BlockMatrix):
- B = unpacked.blocks
- new_B = [
- [factor * B[i, j] for j in range(B.cols)] for i in range(B.rows)]
- return BlockMatrix(new_B)
- return expr
- def bc_matmul(expr):
- if isinstance(expr, MatPow):
- if expr.args[1].is_Integer:
- factor, matrices = (1, [expr.args[0]]*expr.args[1])
- else:
- return expr
- else:
- 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):
- collapse = block_collapse(expr.arg)
- return collapse._eval_transpose()
- def bc_inverse(expr):
- if isinstance(expr.arg, BlockDiagMatrix):
- return expr.inverse()
- 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):
- # See: Inverses of 2x2 Block Matrices, Tzon-Tzer Lu and Sheng-Hua Shiou
- [[A, B],
- [C, D]] = expr.arg.blocks.tolist()
- formula = _choose_2x2_inversion_formula(A, B, C, D)
- if formula != None:
- MI = expr.arg.schur(formula).I
- if formula == 'A':
- AI = A.I
- return BlockMatrix([[AI + AI * B * MI * C * AI, -AI * B * MI], [-MI * C * AI, MI]])
- if formula == 'B':
- BI = B.I
- return BlockMatrix([[-MI * D * BI, MI], [BI + BI * A * MI * D * BI, -BI * A * MI]])
- if formula == 'C':
- CI = C.I
- return BlockMatrix([[-CI * D * MI, CI + CI * D * MI * A * CI], [MI, -MI * A * CI]])
- if formula == 'D':
- DI = D.I
- return BlockMatrix([[MI, -MI * B * DI], [-DI * C * MI, DI + DI * C * MI * B * DI]])
- return expr
- def _choose_2x2_inversion_formula(A, B, C, D):
- """
- Assuming [[A, B], [C, D]] would form a valid square block matrix, find
- which of the classical 2x2 block matrix inversion formulas would be
- best suited.
- Returns 'A', 'B', 'C', 'D' to represent the algorithm involving inversion
- of the given argument or None if the matrix cannot be inverted using
- any of those formulas.
- """
- # Try to find a known invertible matrix. Note that the Schur complement
- # is currently not being considered for this
- A_inv = ask(Q.invertible(A))
- if A_inv == True:
- return 'A'
- B_inv = ask(Q.invertible(B))
- if B_inv == True:
- return 'B'
- C_inv = ask(Q.invertible(C))
- if C_inv == True:
- return 'C'
- D_inv = ask(Q.invertible(D))
- if D_inv == True:
- return 'D'
- # Otherwise try to find a matrix that isn't known to be non-invertible
- if A_inv != False:
- return 'A'
- if B_inv != False:
- return 'B'
- if C_inv != False:
- return 'C'
- if D_inv != False:
- return 'D'
- return None
- 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
- 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(expr):
- """
- Reblock a BlockMatrix so that it has 2x2 blocks of block matrices. If
- possible in such a way that the matrix continues to be invertible using the
- classical 2x2 block inversion formulas.
- """
- if not isinstance(expr, BlockMatrix) or not all(d > 2 for d in expr.blockshape):
- return expr
- BM = BlockMatrix # for brevity's sake
- rowblocks, colblocks = expr.blockshape
- blocks = expr.blocks
- for i in range(1, rowblocks):
- for j in range(1, colblocks):
- # try to split rows at i and cols at j
- A = bc_unpack(BM(blocks[:i, :j]))
- B = bc_unpack(BM(blocks[:i, j:]))
- C = bc_unpack(BM(blocks[i:, :j]))
- D = bc_unpack(BM(blocks[i:, j:]))
- formula = _choose_2x2_inversion_formula(A, B, C, D)
- if formula is not None:
- return BlockMatrix([[A, B], [C, D]])
- # else: nothing worked, just split upper left corner
- return BM([[blocks[0, 0], BM(blocks[0, 1:])],
- [BM(blocks[1:, 0]), BM(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])
|