Lazy Jordan decomposition computation

[sylee957]

I see Jordan decomposition or matrix exponential computation is quite frequently used
But I see a lot of redundancy are made only to make the result an explicit matrix.
For example, matrix exp or log tries to get the diagonal block structure and also spends a lot of time computing the inverse and multiplication for computing the similarity transform.

So, I think that making such results into lazily-evaluated form can greatly improve the performance and usefulness, while making it the most backward-compatible format.

I also find that most of the matrix functional analysis problems come in the form of analyzing the structure of unevaluated jordan form structure rather than the explicit result.
Though I'd say that this lazy construction helps a lot of other decomposition methods, if the decomposition method does not do any special thing other than matrix multiplication or block matrix formulation.

So I'd propose the following changes.

1. Make Matrix.jordan_form return J in BlockDiagMatrix rather than Matrix
2. Make Matrix.exp and Matrix.log return the MatMul than Matrix

And the result should look like

after the diff

diff --git a/sympy/matrices/matrices.py b/sympy/matrices/matrices.py
index 8191f6295f..135583fe54 100644
--- a/sympy/matrices/matrices.py
+++ b/sympy/matrices/matrices.py
@@ -1593,7 +1593,6 @@ def analytic_func(self, f, x):


     def exp(self):
-
         """Return the exponential of a square matrix

         Examples
@@ -1613,21 +1612,16 @@ def exp(self):
                 "Exponentiation is valid only for square matrices")
         try:
             P, J = self.jordan_form()
-            cells = J.get_diag_blocks()
+            cells = J.args
         except MatrixError:
             raise NotImplementedError(
                 "Exponentiation is implemented only for matrices for which the Jordan normal form can be computed")

         blocks = [cell._eval_matrix_exp_jblock() for cell in cells]
-        from sympy.matrices import diag
-        from sympy import re
-        eJ = diag(*blocks)
-        # n = self.rows
-        ret = P.multiply(eJ, dotprodsimp=True).multiply(P.inv(), dotprodsimp=True)
-        if all(value.is_real for value in self.values()):
-            return type(self)(re(ret))
-        else:
-            return type(self)(ret)
+        from .expressions import BlockDiagMatrix, MatMul, Inverse
+        eJ = BlockDiagMatrix(*blocks)
+        return MatMul(P, eJ, Inverse(P))
+

     def _eval_matrix_log_jblock(self):
         """Helper function to compute logarithm of a jordan block.
@@ -1730,7 +1724,7 @@ def log(self, simplify=cancel):
             else:
                 P, J = self.jordan_form()

-            cells = J.get_diag_blocks()
+            cells = J.args
         except MatrixError:
             raise NotImplementedError(
                 "Logarithm is implemented only for matrices for which "
@@ -1739,16 +1733,10 @@ def log(self, simplify=cancel):
         blocks = [
             cell._eval_matrix_log_jblock()
             for cell in cells]
-        from sympy.matrices import diag
-        eJ = diag(*blocks)
-
-        if simplify:
-            ret = simplify(P * eJ * simplify(P.inv()))
-            ret = self.__class__(ret)
-        else:
-            ret = P * eJ * P.inv()
+        from .expressions import BlockDiagMatrix, MatMul, Inverse
+        eJ = BlockDiagMatrix(*blocks)
+        return MatMul(P, eJ, Inverse(P))

-        return ret

     def is_nilpotent(self):
         """Checks if a matrix is nilpotent.
(END)

[oscarbenjamin]

That seems reasonable to me. The documentation should explain how to access the different elements of the the results returned in this way e.g.:

In [9]: B = BlockDiagMatrix(Matrix([[1]]), Matrix([[2, 1], [0, 2]]))                                                                                          

In [10]: B                                                                                                                                                    
Out[10]: 
⎡[1]    𝟘   ⎤
⎢           ⎥
⎢     ⎡2  1⎤⎥
⎢ 𝟘   ⎢    ⎥⎥
⎣     ⎣0  2⎦⎦

In [11]: B.diag                                                                                                                                               
Out[11]: 
⎛     ⎡2  1⎤⎞
⎜[1], ⎢    ⎥⎟
⎝     ⎣0  2⎦⎠

[oscarbenjamin]

Maybe this could also work with matrix symbols like

In [1]: M = MatrixSymbol('M', 2, 2)                                                                                                            

In [2]: JordanForm(M)
JordanVectors(M)*JordanBlocks(M)*JordanVector(M)**-1