Coverage for /usr/lib/python3/dist-packages/sympy/matrices/expressions/blockmatrix.py: 20%

430 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-14 15:55 +0200

1from sympy.assumptions.ask import (Q, ask) 

2from sympy.core import Basic, Add, Mul, S 

3from sympy.core.sympify import _sympify 

4from sympy.functions import adjoint 

5from sympy.functions.elementary.complexes import re, im 

6from sympy.strategies import typed, exhaust, condition, do_one, unpack 

7from sympy.strategies.traverse import bottom_up 

8from sympy.utilities.iterables import is_sequence, sift 

9from sympy.utilities.misc import filldedent 

10 

11from sympy.matrices import Matrix, ShapeError 

12from sympy.matrices.common import NonInvertibleMatrixError 

13from sympy.matrices.expressions.determinant import det, Determinant 

14from sympy.matrices.expressions.inverse import Inverse 

15from sympy.matrices.expressions.matadd import MatAdd 

16from sympy.matrices.expressions.matexpr import MatrixExpr, MatrixElement 

17from sympy.matrices.expressions.matmul import MatMul 

18from sympy.matrices.expressions.matpow import MatPow 

19from sympy.matrices.expressions.slice import MatrixSlice 

20from sympy.matrices.expressions.special import ZeroMatrix, Identity 

21from sympy.matrices.expressions.trace import trace 

22from sympy.matrices.expressions.transpose import Transpose, transpose 

23 

24 

25class BlockMatrix(MatrixExpr): 

26 """A BlockMatrix is a Matrix comprised of other matrices. 

27 

28 The submatrices are stored in a SymPy Matrix object but accessed as part of 

29 a Matrix Expression 

30 

31 >>> from sympy import (MatrixSymbol, BlockMatrix, symbols, 

32 ... Identity, ZeroMatrix, block_collapse) 

33 >>> n,m,l = symbols('n m l') 

34 >>> X = MatrixSymbol('X', n, n) 

35 >>> Y = MatrixSymbol('Y', m, m) 

36 >>> Z = MatrixSymbol('Z', n, m) 

37 >>> B = BlockMatrix([[X, Z], [ZeroMatrix(m,n), Y]]) 

38 >>> print(B) 

39 Matrix([ 

40 [X, Z], 

41 [0, Y]]) 

42 

43 >>> C = BlockMatrix([[Identity(n), Z]]) 

44 >>> print(C) 

45 Matrix([[I, Z]]) 

46 

47 >>> print(block_collapse(C*B)) 

48 Matrix([[X, Z + Z*Y]]) 

49 

50 Some matrices might be comprised of rows of blocks with 

51 the matrices in each row having the same height and the 

52 rows all having the same total number of columns but 

53 not having the same number of columns for each matrix 

54 in each row. In this case, the matrix is not a block 

55 matrix and should be instantiated by Matrix. 

56 

57 >>> from sympy import ones, Matrix 

58 >>> dat = [ 

59 ... [ones(3,2), ones(3,3)*2], 

60 ... [ones(2,3)*3, ones(2,2)*4]] 

61 ... 

62 >>> BlockMatrix(dat) 

63 Traceback (most recent call last): 

64 ... 

65 ValueError: 

66 Although this matrix is comprised of blocks, the blocks do not fill 

67 the matrix in a size-symmetric fashion. To create a full matrix from 

68 these arguments, pass them directly to Matrix. 

69 >>> Matrix(dat) 

70 Matrix([ 

71 [1, 1, 2, 2, 2], 

72 [1, 1, 2, 2, 2], 

73 [1, 1, 2, 2, 2], 

74 [3, 3, 3, 4, 4], 

75 [3, 3, 3, 4, 4]]) 

76 

77 See Also 

78 ======== 

79 sympy.matrices.matrices.MatrixBase.irregular 

80 """ 

81 def __new__(cls, *args, **kwargs): 

82 from sympy.matrices.immutable import ImmutableDenseMatrix 

83 isMat = lambda i: getattr(i, 'is_Matrix', False) 

84 if len(args) != 1 or \ 

85 not is_sequence(args[0]) or \ 

86 len({isMat(r) for r in args[0]}) != 1: 

87 raise ValueError(filldedent(''' 

88 expecting a sequence of 1 or more rows 

89 containing Matrices.''')) 

90 rows = args[0] if args else [] 

91 if not isMat(rows): 

92 if rows and isMat(rows[0]): 

93 rows = [rows] # rows is not list of lists or [] 

94 # regularity check 

95 # same number of matrices in each row 

96 blocky = ok = len({len(r) for r in rows}) == 1 

97 if ok: 

98 # same number of rows for each matrix in a row 

99 for r in rows: 

100 ok = len({i.rows for i in r}) == 1 

101 if not ok: 

102 break 

103 blocky = ok 

104 if ok: 

105 # same number of cols for each matrix in each col 

106 for c in range(len(rows[0])): 

107 ok = len({rows[i][c].cols 

108 for i in range(len(rows))}) == 1 

109 if not ok: 

110 break 

111 if not ok: 

112 # same total cols in each row 

113 ok = len({ 

114 sum([i.cols for i in r]) for r in rows}) == 1 

115 if blocky and ok: 

116 raise ValueError(filldedent(''' 

117 Although this matrix is comprised of blocks, 

118 the blocks do not fill the matrix in a 

119 size-symmetric fashion. To create a full matrix 

120 from these arguments, pass them directly to 

121 Matrix.''')) 

122 raise ValueError(filldedent(''' 

123 When there are not the same number of rows in each 

124 row's matrices or there are not the same number of 

125 total columns in each row, the matrix is not a 

126 block matrix. If this matrix is known to consist of 

127 blocks fully filling a 2-D space then see 

128 Matrix.irregular.''')) 

129 mat = ImmutableDenseMatrix(rows, evaluate=False) 

130 obj = Basic.__new__(cls, mat) 

131 return obj 

132 

133 @property 

134 def shape(self): 

135 numrows = numcols = 0 

136 M = self.blocks 

137 for i in range(M.shape[0]): 

138 numrows += M[i, 0].shape[0] 

139 for i in range(M.shape[1]): 

140 numcols += M[0, i].shape[1] 

141 return (numrows, numcols) 

142 

143 @property 

144 def blockshape(self): 

145 return self.blocks.shape 

146 

147 @property 

148 def blocks(self): 

149 return self.args[0] 

150 

151 @property 

152 def rowblocksizes(self): 

153 return [self.blocks[i, 0].rows for i in range(self.blockshape[0])] 

154 

155 @property 

156 def colblocksizes(self): 

157 return [self.blocks[0, i].cols for i in range(self.blockshape[1])] 

158 

159 def structurally_equal(self, other): 

160 return (isinstance(other, BlockMatrix) 

161 and self.shape == other.shape 

162 and self.blockshape == other.blockshape 

163 and self.rowblocksizes == other.rowblocksizes 

164 and self.colblocksizes == other.colblocksizes) 

165 

166 def _blockmul(self, other): 

167 if (isinstance(other, BlockMatrix) and 

168 self.colblocksizes == other.rowblocksizes): 

169 return BlockMatrix(self.blocks*other.blocks) 

170 

171 return self * other 

172 

173 def _blockadd(self, other): 

174 if (isinstance(other, BlockMatrix) 

175 and self.structurally_equal(other)): 

176 return BlockMatrix(self.blocks + other.blocks) 

177 

178 return self + other 

179 

180 def _eval_transpose(self): 

181 # Flip all the individual matrices 

182 matrices = [transpose(matrix) for matrix in self.blocks] 

183 # Make a copy 

184 M = Matrix(self.blockshape[0], self.blockshape[1], matrices) 

185 # Transpose the block structure 

186 M = M.transpose() 

187 return BlockMatrix(M) 

188 

189 def _eval_adjoint(self): 

190 # Adjoint all the individual matrices 

191 matrices = [adjoint(matrix) for matrix in self.blocks] 

192 # Make a copy 

193 M = Matrix(self.blockshape[0], self.blockshape[1], matrices) 

194 # Transpose the block structure 

195 M = M.transpose() 

196 return BlockMatrix(M) 

197 

198 def _eval_trace(self): 

199 if self.rowblocksizes == self.colblocksizes: 

200 return Add(*[trace(self.blocks[i, i]) 

201 for i in range(self.blockshape[0])]) 

202 raise NotImplementedError( 

203 "Can't perform trace of irregular blockshape") 

204 

205 def _eval_determinant(self): 

206 if self.blockshape == (1, 1): 

207 return det(self.blocks[0, 0]) 

208 if self.blockshape == (2, 2): 

209 [[A, B], 

210 [C, D]] = self.blocks.tolist() 

211 if ask(Q.invertible(A)): 

212 return det(A)*det(D - C*A.I*B) 

213 elif ask(Q.invertible(D)): 

214 return det(D)*det(A - B*D.I*C) 

215 return Determinant(self) 

216 

217 def _eval_as_real_imag(self): 

218 real_matrices = [re(matrix) for matrix in self.blocks] 

219 real_matrices = Matrix(self.blockshape[0], self.blockshape[1], real_matrices) 

220 

221 im_matrices = [im(matrix) for matrix in self.blocks] 

222 im_matrices = Matrix(self.blockshape[0], self.blockshape[1], im_matrices) 

223 

224 return (BlockMatrix(real_matrices), BlockMatrix(im_matrices)) 

225 

226 def transpose(self): 

227 """Return transpose of matrix. 

228 

229 Examples 

230 ======== 

231 

232 >>> from sympy import MatrixSymbol, BlockMatrix, ZeroMatrix 

233 >>> from sympy.abc import m, n 

234 >>> X = MatrixSymbol('X', n, n) 

235 >>> Y = MatrixSymbol('Y', m, m) 

236 >>> Z = MatrixSymbol('Z', n, m) 

237 >>> B = BlockMatrix([[X, Z], [ZeroMatrix(m,n), Y]]) 

238 >>> B.transpose() 

239 Matrix([ 

240 [X.T, 0], 

241 [Z.T, Y.T]]) 

242 >>> _.transpose() 

243 Matrix([ 

244 [X, Z], 

245 [0, Y]]) 

246 """ 

247 return self._eval_transpose() 

248 

249 def schur(self, mat = 'A', generalized = False): 

250 """Return the Schur Complement of the 2x2 BlockMatrix 

251 

252 Parameters 

253 ========== 

254 

255 mat : String, optional 

256 The matrix with respect to which the 

257 Schur Complement is calculated. 'A' is 

258 used by default 

259 

260 generalized : bool, optional 

261 If True, returns the generalized Schur 

262 Component which uses Moore-Penrose Inverse 

263 

264 Examples 

265 ======== 

266 

267 >>> from sympy import symbols, MatrixSymbol, BlockMatrix 

268 >>> m, n = symbols('m n') 

269 >>> A = MatrixSymbol('A', n, n) 

270 >>> B = MatrixSymbol('B', n, m) 

271 >>> C = MatrixSymbol('C', m, n) 

272 >>> D = MatrixSymbol('D', m, m) 

273 >>> X = BlockMatrix([[A, B], [C, D]]) 

274 

275 The default Schur Complement is evaluated with "A" 

276 

277 >>> X.schur() 

278 -C*A**(-1)*B + D 

279 >>> X.schur('D') 

280 A - B*D**(-1)*C 

281 

282 Schur complement with non-invertible matrices is not 

283 defined. Instead, the generalized Schur complement can 

284 be calculated which uses the Moore-Penrose Inverse. To 

285 achieve this, `generalized` must be set to `True` 

286 

287 >>> X.schur('B', generalized=True) 

288 C - D*(B.T*B)**(-1)*B.T*A 

289 >>> X.schur('C', generalized=True) 

290 -A*(C.T*C)**(-1)*C.T*D + B 

291 

292 Returns 

293 ======= 

294 

295 M : Matrix 

296 The Schur Complement Matrix 

297 

298 Raises 

299 ====== 

300 

301 ShapeError 

302 If the block matrix is not a 2x2 matrix 

303 

304 NonInvertibleMatrixError 

305 If given matrix is non-invertible 

306 

307 References 

308 ========== 

309 

310 .. [1] Wikipedia Article on Schur Component : https://en.wikipedia.org/wiki/Schur_complement 

311 

312 See Also 

313 ======== 

314 

315 sympy.matrices.matrices.MatrixBase.pinv 

316 """ 

317 

318 if self.blockshape == (2, 2): 

319 [[A, B], 

320 [C, D]] = self.blocks.tolist() 

321 d={'A' : A, 'B' : B, 'C' : C, 'D' : D} 

322 try: 

323 inv = (d[mat].T*d[mat]).inv()*d[mat].T if generalized else d[mat].inv() 

324 if mat == 'A': 

325 return D - C * inv * B 

326 elif mat == 'B': 

327 return C - D * inv * A 

328 elif mat == 'C': 

329 return B - A * inv * D 

330 elif mat == 'D': 

331 return A - B * inv * C 

332 #For matrices where no sub-matrix is square 

333 return self 

334 except NonInvertibleMatrixError: 

335 raise NonInvertibleMatrixError('The given matrix is not invertible. Please set generalized=True \ 

336 to compute the generalized Schur Complement which uses Moore-Penrose Inverse') 

337 else: 

338 raise ShapeError('Schur Complement can only be calculated for 2x2 block matrices') 

339 

340 def LDUdecomposition(self): 

341 """Returns the Block LDU decomposition of 

342 a 2x2 Block Matrix 

343 

344 Returns 

345 ======= 

346 

347 (L, D, U) : Matrices 

348 L : Lower Diagonal Matrix 

349 D : Diagonal Matrix 

350 U : Upper Diagonal Matrix 

351 

352 Examples 

353 ======== 

354 

355 >>> from sympy import symbols, MatrixSymbol, BlockMatrix, block_collapse 

356 >>> m, n = symbols('m n') 

357 >>> A = MatrixSymbol('A', n, n) 

358 >>> B = MatrixSymbol('B', n, m) 

359 >>> C = MatrixSymbol('C', m, n) 

360 >>> D = MatrixSymbol('D', m, m) 

361 >>> X = BlockMatrix([[A, B], [C, D]]) 

362 >>> L, D, U = X.LDUdecomposition() 

363 >>> block_collapse(L*D*U) 

364 Matrix([ 

365 [A, B], 

366 [C, D]]) 

367 

368 Raises 

369 ====== 

370 

371 ShapeError 

372 If the block matrix is not a 2x2 matrix 

373 

374 NonInvertibleMatrixError 

375 If the matrix "A" is non-invertible 

376 

377 See Also 

378 ======== 

379 sympy.matrices.expressions.blockmatrix.BlockMatrix.UDLdecomposition 

380 sympy.matrices.expressions.blockmatrix.BlockMatrix.LUdecomposition 

381 """ 

382 if self.blockshape == (2,2): 

383 [[A, B], 

384 [C, D]] = self.blocks.tolist() 

385 try: 

386 AI = A.I 

387 except NonInvertibleMatrixError: 

388 raise NonInvertibleMatrixError('Block LDU decomposition cannot be calculated when\ 

389 "A" is singular') 

390 Ip = Identity(B.shape[0]) 

391 Iq = Identity(B.shape[1]) 

392 Z = ZeroMatrix(*B.shape) 

393 L = BlockMatrix([[Ip, Z], [C*AI, Iq]]) 

394 D = BlockDiagMatrix(A, self.schur()) 

395 U = BlockMatrix([[Ip, AI*B],[Z.T, Iq]]) 

396 return L, D, U 

397 else: 

398 raise ShapeError("Block LDU decomposition is supported only for 2x2 block matrices") 

399 

400 def UDLdecomposition(self): 

401 """Returns the Block UDL decomposition of 

402 a 2x2 Block Matrix 

403 

404 Returns 

405 ======= 

406 

407 (U, D, L) : Matrices 

408 U : Upper Diagonal Matrix 

409 D : Diagonal Matrix 

410 L : Lower Diagonal Matrix 

411 

412 Examples 

413 ======== 

414 

415 >>> from sympy import symbols, MatrixSymbol, BlockMatrix, block_collapse 

416 >>> m, n = symbols('m n') 

417 >>> A = MatrixSymbol('A', n, n) 

418 >>> B = MatrixSymbol('B', n, m) 

419 >>> C = MatrixSymbol('C', m, n) 

420 >>> D = MatrixSymbol('D', m, m) 

421 >>> X = BlockMatrix([[A, B], [C, D]]) 

422 >>> U, D, L = X.UDLdecomposition() 

423 >>> block_collapse(U*D*L) 

424 Matrix([ 

425 [A, B], 

426 [C, D]]) 

427 

428 Raises 

429 ====== 

430 

431 ShapeError 

432 If the block matrix is not a 2x2 matrix 

433 

434 NonInvertibleMatrixError 

435 If the matrix "D" is non-invertible 

436 

437 See Also 

438 ======== 

439 sympy.matrices.expressions.blockmatrix.BlockMatrix.LDUdecomposition 

440 sympy.matrices.expressions.blockmatrix.BlockMatrix.LUdecomposition 

441 """ 

442 if self.blockshape == (2,2): 

443 [[A, B], 

444 [C, D]] = self.blocks.tolist() 

445 try: 

446 DI = D.I 

447 except NonInvertibleMatrixError: 

448 raise NonInvertibleMatrixError('Block UDL decomposition cannot be calculated when\ 

449 "D" is singular') 

450 Ip = Identity(A.shape[0]) 

451 Iq = Identity(B.shape[1]) 

452 Z = ZeroMatrix(*B.shape) 

453 U = BlockMatrix([[Ip, B*DI], [Z.T, Iq]]) 

454 D = BlockDiagMatrix(self.schur('D'), D) 

455 L = BlockMatrix([[Ip, Z],[DI*C, Iq]]) 

456 return U, D, L 

457 else: 

458 raise ShapeError("Block UDL decomposition is supported only for 2x2 block matrices") 

459 

460 def LUdecomposition(self): 

461 """Returns the Block LU decomposition of 

462 a 2x2 Block Matrix 

463 

464 Returns 

465 ======= 

466 

467 (L, U) : Matrices 

468 L : Lower Diagonal Matrix 

469 U : Upper Diagonal Matrix 

470 

471 Examples 

472 ======== 

473 

474 >>> from sympy import symbols, MatrixSymbol, BlockMatrix, block_collapse 

475 >>> m, n = symbols('m n') 

476 >>> A = MatrixSymbol('A', n, n) 

477 >>> B = MatrixSymbol('B', n, m) 

478 >>> C = MatrixSymbol('C', m, n) 

479 >>> D = MatrixSymbol('D', m, m) 

480 >>> X = BlockMatrix([[A, B], [C, D]]) 

481 >>> L, U = X.LUdecomposition() 

482 >>> block_collapse(L*U) 

483 Matrix([ 

484 [A, B], 

485 [C, D]]) 

486 

487 Raises 

488 ====== 

489 

490 ShapeError 

491 If the block matrix is not a 2x2 matrix 

492 

493 NonInvertibleMatrixError 

494 If the matrix "A" is non-invertible 

495 

496 See Also 

497 ======== 

498 sympy.matrices.expressions.blockmatrix.BlockMatrix.UDLdecomposition 

499 sympy.matrices.expressions.blockmatrix.BlockMatrix.LDUdecomposition 

500 """ 

501 if self.blockshape == (2,2): 

502 [[A, B], 

503 [C, D]] = self.blocks.tolist() 

504 try: 

505 A = A**0.5 

506 AI = A.I 

507 except NonInvertibleMatrixError: 

508 raise NonInvertibleMatrixError('Block LU decomposition cannot be calculated when\ 

509 "A" is singular') 

510 Z = ZeroMatrix(*B.shape) 

511 Q = self.schur()**0.5 

512 L = BlockMatrix([[A, Z], [C*AI, Q]]) 

513 U = BlockMatrix([[A, AI*B],[Z.T, Q]]) 

514 return L, U 

515 else: 

516 raise ShapeError("Block LU decomposition is supported only for 2x2 block matrices") 

517 

518 def _entry(self, i, j, **kwargs): 

519 # Find row entry 

520 orig_i, orig_j = i, j 

521 for row_block, numrows in enumerate(self.rowblocksizes): 

522 cmp = i < numrows 

523 if cmp == True: 

524 break 

525 elif cmp == False: 

526 i -= numrows 

527 elif row_block < self.blockshape[0] - 1: 

528 # Can't tell which block and it's not the last one, return unevaluated 

529 return MatrixElement(self, orig_i, orig_j) 

530 for col_block, numcols in enumerate(self.colblocksizes): 

531 cmp = j < numcols 

532 if cmp == True: 

533 break 

534 elif cmp == False: 

535 j -= numcols 

536 elif col_block < self.blockshape[1] - 1: 

537 return MatrixElement(self, orig_i, orig_j) 

538 return self.blocks[row_block, col_block][i, j] 

539 

540 @property 

541 def is_Identity(self): 

542 if self.blockshape[0] != self.blockshape[1]: 

543 return False 

544 for i in range(self.blockshape[0]): 

545 for j in range(self.blockshape[1]): 

546 if i==j and not self.blocks[i, j].is_Identity: 

547 return False 

548 if i!=j and not self.blocks[i, j].is_ZeroMatrix: 

549 return False 

550 return True 

551 

552 @property 

553 def is_structurally_symmetric(self): 

554 return self.rowblocksizes == self.colblocksizes 

555 

556 def equals(self, other): 

557 if self == other: 

558 return True 

559 if (isinstance(other, BlockMatrix) and self.blocks == other.blocks): 

560 return True 

561 return super().equals(other) 

562 

563 

564class BlockDiagMatrix(BlockMatrix): 

565 """A sparse matrix with block matrices along its diagonals 

566 

567 Examples 

568 ======== 

569 

570 >>> from sympy import MatrixSymbol, BlockDiagMatrix, symbols 

571 >>> n, m, l = symbols('n m l') 

572 >>> X = MatrixSymbol('X', n, n) 

573 >>> Y = MatrixSymbol('Y', m, m) 

574 >>> BlockDiagMatrix(X, Y) 

575 Matrix([ 

576 [X, 0], 

577 [0, Y]]) 

578 

579 Notes 

580 ===== 

581 

582 If you want to get the individual diagonal blocks, use 

583 :meth:`get_diag_blocks`. 

584 

585 See Also 

586 ======== 

587 

588 sympy.matrices.dense.diag 

589 """ 

590 def __new__(cls, *mats): 

591 return Basic.__new__(BlockDiagMatrix, *[_sympify(m) for m in mats]) 

592 

593 @property 

594 def diag(self): 

595 return self.args 

596 

597 @property 

598 def blocks(self): 

599 from sympy.matrices.immutable import ImmutableDenseMatrix 

600 mats = self.args 

601 data = [[mats[i] if i == j else ZeroMatrix(mats[i].rows, mats[j].cols) 

602 for j in range(len(mats))] 

603 for i in range(len(mats))] 

604 return ImmutableDenseMatrix(data, evaluate=False) 

605 

606 @property 

607 def shape(self): 

608 return (sum(block.rows for block in self.args), 

609 sum(block.cols for block in self.args)) 

610 

611 @property 

612 def blockshape(self): 

613 n = len(self.args) 

614 return (n, n) 

615 

616 @property 

617 def rowblocksizes(self): 

618 return [block.rows for block in self.args] 

619 

620 @property 

621 def colblocksizes(self): 

622 return [block.cols for block in self.args] 

623 

624 def _all_square_blocks(self): 

625 """Returns true if all blocks are square""" 

626 return all(mat.is_square for mat in self.args) 

627 

628 def _eval_determinant(self): 

629 if self._all_square_blocks(): 

630 return Mul(*[det(mat) for mat in self.args]) 

631 # At least one block is non-square. Since the entire matrix must be square we know there must 

632 # be at least two blocks in this matrix, in which case the entire matrix is necessarily rank-deficient 

633 return S.Zero 

634 

635 def _eval_inverse(self, expand='ignored'): 

636 if self._all_square_blocks(): 

637 return BlockDiagMatrix(*[mat.inverse() for mat in self.args]) 

638 # See comment in _eval_determinant() 

639 raise NonInvertibleMatrixError('Matrix det == 0; not invertible.') 

640 

641 def _eval_transpose(self): 

642 return BlockDiagMatrix(*[mat.transpose() for mat in self.args]) 

643 

644 def _blockmul(self, other): 

645 if (isinstance(other, BlockDiagMatrix) and 

646 self.colblocksizes == other.rowblocksizes): 

647 return BlockDiagMatrix(*[a*b for a, b in zip(self.args, other.args)]) 

648 else: 

649 return BlockMatrix._blockmul(self, other) 

650 

651 def _blockadd(self, other): 

652 if (isinstance(other, BlockDiagMatrix) and 

653 self.blockshape == other.blockshape and 

654 self.rowblocksizes == other.rowblocksizes and 

655 self.colblocksizes == other.colblocksizes): 

656 return BlockDiagMatrix(*[a + b for a, b in zip(self.args, other.args)]) 

657 else: 

658 return BlockMatrix._blockadd(self, other) 

659 

660 def get_diag_blocks(self): 

661 """Return the list of diagonal blocks of the matrix. 

662 

663 Examples 

664 ======== 

665 

666 >>> from sympy import BlockDiagMatrix, Matrix 

667 

668 >>> A = Matrix([[1, 2], [3, 4]]) 

669 >>> B = Matrix([[5, 6], [7, 8]]) 

670 >>> M = BlockDiagMatrix(A, B) 

671 

672 How to get diagonal blocks from the block diagonal matrix: 

673 

674 >>> diag_blocks = M.get_diag_blocks() 

675 >>> diag_blocks[0] 

676 Matrix([ 

677 [1, 2], 

678 [3, 4]]) 

679 >>> diag_blocks[1] 

680 Matrix([ 

681 [5, 6], 

682 [7, 8]]) 

683 """ 

684 return self.args 

685 

686 

687def block_collapse(expr): 

688 """Evaluates a block matrix expression 

689 

690 >>> from sympy import MatrixSymbol, BlockMatrix, symbols, Identity, ZeroMatrix, block_collapse 

691 >>> n,m,l = symbols('n m l') 

692 >>> X = MatrixSymbol('X', n, n) 

693 >>> Y = MatrixSymbol('Y', m, m) 

694 >>> Z = MatrixSymbol('Z', n, m) 

695 >>> B = BlockMatrix([[X, Z], [ZeroMatrix(m, n), Y]]) 

696 >>> print(B) 

697 Matrix([ 

698 [X, Z], 

699 [0, Y]]) 

700 

701 >>> C = BlockMatrix([[Identity(n), Z]]) 

702 >>> print(C) 

703 Matrix([[I, Z]]) 

704 

705 >>> print(block_collapse(C*B)) 

706 Matrix([[X, Z + Z*Y]]) 

707 """ 

708 from sympy.strategies.util import expr_fns 

709 

710 hasbm = lambda expr: isinstance(expr, MatrixExpr) and expr.has(BlockMatrix) 

711 

712 conditioned_rl = condition( 

713 hasbm, 

714 typed( 

715 {MatAdd: do_one(bc_matadd, bc_block_plus_ident), 

716 MatMul: do_one(bc_matmul, bc_dist), 

717 MatPow: bc_matmul, 

718 Transpose: bc_transpose, 

719 Inverse: bc_inverse, 

720 BlockMatrix: do_one(bc_unpack, deblock)} 

721 ) 

722 ) 

723 

724 rule = exhaust( 

725 bottom_up( 

726 exhaust(conditioned_rl), 

727 fns=expr_fns 

728 ) 

729 ) 

730 

731 result = rule(expr) 

732 doit = getattr(result, 'doit', None) 

733 if doit is not None: 

734 return doit() 

735 else: 

736 return result 

737 

738def bc_unpack(expr): 

739 if expr.blockshape == (1, 1): 

740 return expr.blocks[0, 0] 

741 return expr 

742 

743def bc_matadd(expr): 

744 args = sift(expr.args, lambda M: isinstance(M, BlockMatrix)) 

745 blocks = args[True] 

746 if not blocks: 

747 return expr 

748 

749 nonblocks = args[False] 

750 block = blocks[0] 

751 for b in blocks[1:]: 

752 block = block._blockadd(b) 

753 if nonblocks: 

754 return MatAdd(*nonblocks) + block 

755 else: 

756 return block 

757 

758def bc_block_plus_ident(expr): 

759 idents = [arg for arg in expr.args if arg.is_Identity] 

760 if not idents: 

761 return expr 

762 

763 blocks = [arg for arg in expr.args if isinstance(arg, BlockMatrix)] 

764 if (blocks and all(b.structurally_equal(blocks[0]) for b in blocks) 

765 and blocks[0].is_structurally_symmetric): 

766 block_id = BlockDiagMatrix(*[Identity(k) 

767 for k in blocks[0].rowblocksizes]) 

768 rest = [arg for arg in expr.args if not arg.is_Identity and not isinstance(arg, BlockMatrix)] 

769 return MatAdd(block_id * len(idents), *blocks, *rest).doit() 

770 

771 return expr 

772 

773def bc_dist(expr): 

774 """ Turn a*[X, Y] into [a*X, a*Y] """ 

775 factor, mat = expr.as_coeff_mmul() 

776 if factor == 1: 

777 return expr 

778 

779 unpacked = unpack(mat) 

780 

781 if isinstance(unpacked, BlockDiagMatrix): 

782 B = unpacked.diag 

783 new_B = [factor * mat for mat in B] 

784 return BlockDiagMatrix(*new_B) 

785 elif isinstance(unpacked, BlockMatrix): 

786 B = unpacked.blocks 

787 new_B = [ 

788 [factor * B[i, j] for j in range(B.cols)] for i in range(B.rows)] 

789 return BlockMatrix(new_B) 

790 return expr 

791 

792 

793def bc_matmul(expr): 

794 if isinstance(expr, MatPow): 

795 if expr.args[1].is_Integer: 

796 factor, matrices = (1, [expr.args[0]]*expr.args[1]) 

797 else: 

798 return expr 

799 else: 

800 factor, matrices = expr.as_coeff_matrices() 

801 

802 i = 0 

803 while (i+1 < len(matrices)): 

804 A, B = matrices[i:i+2] 

805 if isinstance(A, BlockMatrix) and isinstance(B, BlockMatrix): 

806 matrices[i] = A._blockmul(B) 

807 matrices.pop(i+1) 

808 elif isinstance(A, BlockMatrix): 

809 matrices[i] = A._blockmul(BlockMatrix([[B]])) 

810 matrices.pop(i+1) 

811 elif isinstance(B, BlockMatrix): 

812 matrices[i] = BlockMatrix([[A]])._blockmul(B) 

813 matrices.pop(i+1) 

814 else: 

815 i+=1 

816 return MatMul(factor, *matrices).doit() 

817 

818def bc_transpose(expr): 

819 collapse = block_collapse(expr.arg) 

820 return collapse._eval_transpose() 

821 

822 

823def bc_inverse(expr): 

824 if isinstance(expr.arg, BlockDiagMatrix): 

825 return expr.inverse() 

826 

827 expr2 = blockinverse_1x1(expr) 

828 if expr != expr2: 

829 return expr2 

830 return blockinverse_2x2(Inverse(reblock_2x2(expr.arg))) 

831 

832def blockinverse_1x1(expr): 

833 if isinstance(expr.arg, BlockMatrix) and expr.arg.blockshape == (1, 1): 

834 mat = Matrix([[expr.arg.blocks[0].inverse()]]) 

835 return BlockMatrix(mat) 

836 return expr 

837 

838 

839def blockinverse_2x2(expr): 

840 if isinstance(expr.arg, BlockMatrix) and expr.arg.blockshape == (2, 2): 

841 # See: Inverses of 2x2 Block Matrices, Tzon-Tzer Lu and Sheng-Hua Shiou 

842 [[A, B], 

843 [C, D]] = expr.arg.blocks.tolist() 

844 

845 formula = _choose_2x2_inversion_formula(A, B, C, D) 

846 if formula != None: 

847 MI = expr.arg.schur(formula).I 

848 if formula == 'A': 

849 AI = A.I 

850 return BlockMatrix([[AI + AI * B * MI * C * AI, -AI * B * MI], [-MI * C * AI, MI]]) 

851 if formula == 'B': 

852 BI = B.I 

853 return BlockMatrix([[-MI * D * BI, MI], [BI + BI * A * MI * D * BI, -BI * A * MI]]) 

854 if formula == 'C': 

855 CI = C.I 

856 return BlockMatrix([[-CI * D * MI, CI + CI * D * MI * A * CI], [MI, -MI * A * CI]]) 

857 if formula == 'D': 

858 DI = D.I 

859 return BlockMatrix([[MI, -MI * B * DI], [-DI * C * MI, DI + DI * C * MI * B * DI]]) 

860 

861 return expr 

862 

863 

864def _choose_2x2_inversion_formula(A, B, C, D): 

865 """ 

866 Assuming [[A, B], [C, D]] would form a valid square block matrix, find 

867 which of the classical 2x2 block matrix inversion formulas would be 

868 best suited. 

869 

870 Returns 'A', 'B', 'C', 'D' to represent the algorithm involving inversion 

871 of the given argument or None if the matrix cannot be inverted using 

872 any of those formulas. 

873 """ 

874 # Try to find a known invertible matrix. Note that the Schur complement 

875 # is currently not being considered for this 

876 A_inv = ask(Q.invertible(A)) 

877 if A_inv == True: 

878 return 'A' 

879 B_inv = ask(Q.invertible(B)) 

880 if B_inv == True: 

881 return 'B' 

882 C_inv = ask(Q.invertible(C)) 

883 if C_inv == True: 

884 return 'C' 

885 D_inv = ask(Q.invertible(D)) 

886 if D_inv == True: 

887 return 'D' 

888 # Otherwise try to find a matrix that isn't known to be non-invertible 

889 if A_inv != False: 

890 return 'A' 

891 if B_inv != False: 

892 return 'B' 

893 if C_inv != False: 

894 return 'C' 

895 if D_inv != False: 

896 return 'D' 

897 return None 

898 

899 

900def deblock(B): 

901 """ Flatten a BlockMatrix of BlockMatrices """ 

902 if not isinstance(B, BlockMatrix) or not B.blocks.has(BlockMatrix): 

903 return B 

904 wrap = lambda x: x if isinstance(x, BlockMatrix) else BlockMatrix([[x]]) 

905 bb = B.blocks.applyfunc(wrap) # everything is a block 

906 

907 try: 

908 MM = Matrix(0, sum(bb[0, i].blocks.shape[1] for i in range(bb.shape[1])), []) 

909 for row in range(0, bb.shape[0]): 

910 M = Matrix(bb[row, 0].blocks) 

911 for col in range(1, bb.shape[1]): 

912 M = M.row_join(bb[row, col].blocks) 

913 MM = MM.col_join(M) 

914 

915 return BlockMatrix(MM) 

916 except ShapeError: 

917 return B 

918 

919 

920def reblock_2x2(expr): 

921 """ 

922 Reblock a BlockMatrix so that it has 2x2 blocks of block matrices. If 

923 possible in such a way that the matrix continues to be invertible using the 

924 classical 2x2 block inversion formulas. 

925 """ 

926 if not isinstance(expr, BlockMatrix) or not all(d > 2 for d in expr.blockshape): 

927 return expr 

928 

929 BM = BlockMatrix # for brevity's sake 

930 rowblocks, colblocks = expr.blockshape 

931 blocks = expr.blocks 

932 for i in range(1, rowblocks): 

933 for j in range(1, colblocks): 

934 # try to split rows at i and cols at j 

935 A = bc_unpack(BM(blocks[:i, :j])) 

936 B = bc_unpack(BM(blocks[:i, j:])) 

937 C = bc_unpack(BM(blocks[i:, :j])) 

938 D = bc_unpack(BM(blocks[i:, j:])) 

939 

940 formula = _choose_2x2_inversion_formula(A, B, C, D) 

941 if formula is not None: 

942 return BlockMatrix([[A, B], [C, D]]) 

943 

944 # else: nothing worked, just split upper left corner 

945 return BM([[blocks[0, 0], BM(blocks[0, 1:])], 

946 [BM(blocks[1:, 0]), BM(blocks[1:, 1:])]]) 

947 

948 

949def bounds(sizes): 

950 """ Convert sequence of numbers into pairs of low-high pairs 

951 

952 >>> from sympy.matrices.expressions.blockmatrix import bounds 

953 >>> bounds((1, 10, 50)) 

954 [(0, 1), (1, 11), (11, 61)] 

955 """ 

956 low = 0 

957 rv = [] 

958 for size in sizes: 

959 rv.append((low, low + size)) 

960 low += size 

961 return rv 

962 

963def blockcut(expr, rowsizes, colsizes): 

964 """ Cut a matrix expression into Blocks 

965 

966 >>> from sympy import ImmutableMatrix, blockcut 

967 >>> M = ImmutableMatrix(4, 4, range(16)) 

968 >>> B = blockcut(M, (1, 3), (1, 3)) 

969 >>> type(B).__name__ 

970 'BlockMatrix' 

971 >>> ImmutableMatrix(B.blocks[0, 1]) 

972 Matrix([[1, 2, 3]]) 

973 """ 

974 

975 rowbounds = bounds(rowsizes) 

976 colbounds = bounds(colsizes) 

977 return BlockMatrix([[MatrixSlice(expr, rowbound, colbound) 

978 for colbound in colbounds] 

979 for rowbound in rowbounds])