Merge pull request sympy#19215 from jlherren/improve-2x2-block-inversion

sylee957 · web-flow · commit 226b942f0748 · 2020-05-01T11:00:53.000+09:00
Improve 2x2 block matrix inversion with off-diagonal invertible blocks
diff --git a/sympy/assumptions/handlers/matrices.py b/sympy/assumptions/handlers/matrices.py
@@ -155,6 +155,53 @@ def MatrixSlice(expr, assumptions):
         else:
             return ask(Q.invertible(expr.parent), assumptions)
 
+    @staticmethod
+    def MatrixBase(expr, assumptions):
+        if not expr.is_square:
+            return False
+        return expr.rank() == expr.rows
+
+    @staticmethod
+    def MatrixExpr(expr, assumptions):
+        if not expr.is_square:
+            return False
+        return None
+
+    @staticmethod
+    def BlockMatrix(expr, assumptions):
+        from sympy.matrices.expressions.blockmatrix import reblock_2x2
+        if not expr.is_square:
+            return False
+        if expr.blockshape == (1, 1):
+            return ask(Q.invertible(expr.blocks[0, 0]), assumptions)
+        expr = reblock_2x2(expr)
+        if expr.blockshape == (2, 2):
+            [[A, B], [C, D]] = expr.blocks.tolist()
+            if ask(Q.invertible(A), assumptions) == True:
+                invertible = ask(Q.invertible(D - C * A.I * B), assumptions)
+                if invertible is not None:
+                    return invertible
+            if ask(Q.invertible(B), assumptions) == True:
+                invertible = ask(Q.invertible(C - D * B.I * A), assumptions)
+                if invertible is not None:
+                    return invertible
+            if ask(Q.invertible(C), assumptions) == True:
+                invertible = ask(Q.invertible(B - A * C.I * D), assumptions)
+                if invertible is not None:
+                    return invertible
+            if ask(Q.invertible(D), assumptions) == True:
+                invertible = ask(Q.invertible(A - B * D.I * C), assumptions)
+                if invertible is not None:
+                    return invertible
+        return None
+
+    @staticmethod
+    def BlockDiagMatrix(expr, assumptions):
+        if expr.rowblocksizes != expr.colblocksizes:
+            return None
+        return fuzzy_and([ask(Q.invertible(a), assumptions) for a in expr.diag])
+
+
 class AskOrthogonalHandler(CommonHandler):
     """
     Handler for key 'orthogonal'
diff --git a/sympy/assumptions/tests/test_matrices.py b/sympy/assumptions/tests/test_matrices.py
@@ -1,6 +1,7 @@
 from sympy import Q, ask, Symbol, DiagMatrix, DiagonalMatrix
+from sympy.matrices.dense import Matrix
 from sympy.matrices.expressions import (MatrixSymbol, Identity, ZeroMatrix,
-        OneMatrix, Trace, MatrixSlice, Determinant)
+        OneMatrix, Trace, MatrixSlice, Determinant, BlockMatrix, BlockDiagMatrix)
 from sympy.matrices.expressions.factorizations import LofLU
 from sympy.testing.pytest import XFAIL
 
@@ -43,6 +44,40 @@ def test_invertible_fullrank():
     assert ask(Q.invertible(X), Q.fullrank(X)) is True
 
 
+def test_invertible_BlockMatrix():
+    assert ask(Q.invertible(BlockMatrix([Identity(3)]))) == True
+    assert ask(Q.invertible(BlockMatrix([ZeroMatrix(3, 3)]))) == False
+
+    X = Matrix([[1, 2, 3], [3, 5, 4]])
+    Y = Matrix([[4, 2, 7], [2, 3, 5]])
+    # non-invertible A block
+    assert ask(Q.invertible(BlockMatrix([
+        [Matrix.ones(3, 3), Y.T],
+        [X, Matrix.eye(2)],
+    ]))) == True
+    # non-invertible B block
+    assert ask(Q.invertible(BlockMatrix([
+        [Y.T, Matrix.ones(3, 3)],
+        [Matrix.eye(2), X],
+    ]))) == True
+    # non-invertible C block
+    assert ask(Q.invertible(BlockMatrix([
+        [X, Matrix.eye(2)],
+        [Matrix.ones(3, 3), Y.T],
+    ]))) == True
+    # non-invertible D block
+    assert ask(Q.invertible(BlockMatrix([
+        [Matrix.eye(2), X],
+        [Y.T, Matrix.ones(3, 3)],
+    ]))) == True
+
+
+def test_invertible_BlockDiagMatrix():
+    assert ask(Q.invertible(BlockDiagMatrix(Identity(3), Identity(5)))) == True
+    assert ask(Q.invertible(BlockDiagMatrix(ZeroMatrix(3, 3), Identity(5)))) == False
+    assert ask(Q.invertible(BlockDiagMatrix(Identity(3), OneMatrix(5, 5)))) == False
+
+
 def test_symmetric():
     assert ask(Q.symmetric(X), Q.symmetric(X))
     assert ask(Q.symmetric(X*Z), Q.symmetric(X)) is None
@@ -236,7 +271,6 @@ def test_matrix_element_sets():
 
 
 def test_matrix_element_sets_slices_blocks():
-    from sympy.matrices.expressions import BlockMatrix
     X = MatrixSymbol('X', 4, 4)
     assert ask(Q.integer_elements(X[:, 3]), Q.integer_elements(X))
     assert ask(Q.integer_elements(BlockMatrix([[X], [X]])),
diff --git a/sympy/matrices/expressions/blockmatrix.py b/sympy/matrices/expressions/blockmatrix.py
@@ -556,38 +556,71 @@ def blockinverse_1x1(expr):
         return BlockMatrix(mat)
     return expr
 
+
 def blockinverse_2x2(expr):
-    if isinstance(expr.arg, BlockMatrix) and expr.arg.blockshape == (2, 2) and expr.arg.blocks[0, 0].is_square:
-        # Cite: The Matrix Cookbook Section 9.1.3
+    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()
 
-        # Use one or the other formula, depending on whether A or D is known to be invertible or at least not known
-        # to not be invertible.  Note that invertAbility of the other expressions M is not checked.
-        A_invertible = ask(Q.invertible(A))
-        D_invertible = ask(Q.invertible(D))
-        if A_invertible == True:
-            invert_A = True
-        elif D_invertible == True:
-            invert_A = False
-        elif A_invertible != False:
-            invert_A = True
-        elif D_invertible != False:
-            invert_A = False
-        else:
-            invert_A = True
+        formula = _choose_2x2_inversion_formula(A, B, C, D)
 
-        if invert_A:
+        if formula == 'A':
             AI = A.I
             MI = (D - C * AI * B).I
             return BlockMatrix([[AI + AI * B * MI * C * AI, -AI * B * MI], [-MI * C * AI, MI]])
-        else:
+        if formula == 'B':
+            BI = B.I
+            MI = (C - D * BI * A).I
+            return BlockMatrix([[-MI * D * BI, MI], [BI + BI * A * MI * D * BI, -BI * A * MI]])
+        if formula == 'C':
+            CI = C.I
+            MI = (B - A * CI * D).I
+            return BlockMatrix([[-CI * D * MI, CI + CI * D * MI * A * CI], [MI, -MI * A * CI]])
+        if formula == 'D':
             DI = D.I
             MI = (A - B * DI * C).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):
@@ -609,15 +642,33 @@ def deblock(B):
         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
+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
-    return BM([[   B.blocks[0,  0],  BM(B.blocks[0,  1:])],
-               [BM(B.blocks[1:, 0]), BM(B.blocks[1:, 1:])]])
+    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):
diff --git a/sympy/matrices/expressions/tests/test_blockmatrix.py b/sympy/matrices/expressions/tests/test_blockmatrix.py
@@ -1,11 +1,11 @@
 from sympy import Trace
-from sympy.testing.pytest import raises
+from sympy.testing.pytest import raises, slow
 from sympy.matrices.expressions.blockmatrix import (
     block_collapse, bc_matmul, bc_block_plus_ident, BlockDiagMatrix,
     BlockMatrix, bc_dist, bc_matadd, bc_transpose, bc_inverse,
     blockcut, reblock_2x2, deblock)
 from sympy.matrices.expressions import (MatrixSymbol, Identity,
-        Inverse, trace, Transpose, det, ZeroMatrix, OneMatrix)
+        Inverse, trace, Transpose, det, ZeroMatrix)
 from sympy.matrices.common import NonInvertibleMatrixError
 from sympy.matrices import (
     Matrix, ImmutableMatrix, ImmutableSparseMatrix)
@@ -161,35 +161,98 @@ def test_squareBlockMatrix():
     Z = BlockMatrix([[Identity(n), B], [C, D]])
     assert not Z.is_Identity
 
-def test_BlockMatrix_inverse():
+
+def test_BlockMatrix_2x2_inverse_symbolic():
     A = MatrixSymbol('A', n, m)
-    B = MatrixSymbol('B', n, n)
-    C = MatrixSymbol('C', m, m)
-    D = MatrixSymbol('D', m, n)
+    B = MatrixSymbol('B', n, k - m)
+    C = MatrixSymbol('C', k - n, m)
+    D = MatrixSymbol('D', k - n, k - m)
     X = BlockMatrix([[A, B], [C, D]])
-    assert X.is_square
-    assert isinstance(block_collapse(X.inverse()), Inverse)  # Can't inverse when A, D aren't square
+    assert X.is_square and X.shape == (k, k)
+    assert isinstance(block_collapse(X.I), Inverse)  # Can't invert when none of the blocks is square
 
-    # test code path for non-invertible D matrix
+    # test code path where only A is invertible
     A = MatrixSymbol('A', n, n)
     B = MatrixSymbol('B', n, m)
     C = MatrixSymbol('C', m, n)
-    D = OneMatrix(m, m)
+    D = ZeroMatrix(m, m)
     X = BlockMatrix([[A, B], [C, D]])
     assert block_collapse(X.inverse()) == BlockMatrix([
         [A.I + A.I * B * (D - C * A.I * B).I * C * A.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],
     ])
 
-    # test code path for non-invertible A matrix
-    A = OneMatrix(n, n)
+    # test code path where only B is invertible
+    A = MatrixSymbol('A', n, m)
+    B = MatrixSymbol('B', n, n)
+    C = ZeroMatrix(m, m)
+    D = MatrixSymbol('D', m, n)
+    X = BlockMatrix([[A, B], [C, D]])
+    assert block_collapse(X.inverse()) == BlockMatrix(([
+        [-(C - D * B.I * A).I * D * B.I, (C - D * B.I * A).I],
+        [B.I + B.I * A * (C - D * B.I * A).I * D * B.I, -B.I * A * (C - D * B.I * A).I],
+    ]))
+
+    # test code path where only C is invertible
+    A = MatrixSymbol('A', n, m)
+    B = ZeroMatrix(n, n)
+    C = MatrixSymbol('C', m, m)
+    D = MatrixSymbol('D', m, n)
+    X = BlockMatrix([[A, B], [C, D]])
+    assert block_collapse(X.inverse()) == BlockMatrix([
+        [-C.I * D * (B - A * C.I * D).I, C.I + C.I * D * (B - A * C.I * D).I * A * C.I],
+        [(B - A * C.I * D).I, -(B - A * C.I * D).I * A * C.I],
+    ])
+
+    # test code path where only D is invertible
+    A = ZeroMatrix(n, n)
+    B = MatrixSymbol('B', n, m)
+    C = MatrixSymbol('C', m, n)
     D = MatrixSymbol('D', m, m)
     X = BlockMatrix([[A, B], [C, D]])
     assert block_collapse(X.inverse()) == BlockMatrix([
         [(A - B * D.I * C).I, -(A - B * D.I * C).I * B * D.I],
         [-D.I * C * (A - B * D.I * C).I, D.I + D.I * C * (A - B * D.I * C).I * B * D.I],
     ])
 
+
+def test_BlockMatrix_2x2_inverse_numeric():
+    """Test 2x2 block matrix inversion numerically for all 4 formulas"""
+    M = Matrix([[1, 2], [3, 4]])
+    # rank deficient matrices that have full rank when two of them combined
+    D1 = Matrix([[1, 2], [2, 4]])
+    D2 = Matrix([[1, 3], [3, 9]])
+    D3 = Matrix([[1, 4], [4, 16]])
+    assert D1.rank() == D2.rank() == D3.rank() == 1
+    assert (D1 + D2).rank() == (D2 + D3).rank() == (D3 + D1).rank() == 2
+
+    # Only A is invertible
+    K = BlockMatrix([[M, D1], [D2, D3]])
+    assert block_collapse(K.inv()).as_explicit() == K.as_explicit().inv()
+    # Only B is invertible
+    K = BlockMatrix([[D1, M], [D2, D3]])
+    assert block_collapse(K.inv()).as_explicit() == K.as_explicit().inv()
+    # Only C is invertible
+    K = BlockMatrix([[D1, D2], [M, D3]])
+    assert block_collapse(K.inv()).as_explicit() == K.as_explicit().inv()
+    # Only D is invertible
+    K = BlockMatrix([[D1, D2], [D3, M]])
+    assert block_collapse(K.inv()).as_explicit() == K.as_explicit().inv()
+
+
+@slow
+def test_BlockMatrix_3x3_symbolic():
+    # Only test one of these, instead of all permutations, because it's slow
+    rowblocksizes = (n, m, k)
+    colblocksizes = (m, k, n)
+    K = BlockMatrix([
+        [MatrixSymbol('M%s%s' % (rows, cols), rows, cols) for cols in colblocksizes]
+        for rows in rowblocksizes
+    ])
+    collapse = block_collapse(K.I)
+    assert isinstance(collapse, BlockMatrix)
+
+
 def test_BlockDiagMatrix():
     A = MatrixSymbol('A', n, n)
     B = MatrixSymbol('B', m, m)
