fixing matrix derivative bugs

Upabjojr · Upabjojr · commit 0a09bf9e0403 · 2019-03-23T21:53:58.000+01:00
diff --git a/sympy/core/expr.py b/sympy/core/expr.py
@@ -3420,6 +3420,10 @@ def round(self, p=0):
                 allow += 1
             return Float(rv, allow)
 
+    def _eval_derivative_matrix_lines(self, x):
+        from sympy.matrices.expressions.matexpr import _LeftRightArgs
+        return [_LeftRightArgs(S.One, S.One, higher=self._eval_derivative(x))]
+
 
 class AtomicExpr(Atom, Expr):
     """
@@ -3440,9 +3444,9 @@ def _eval_derivative(self, s):
 
     def _eval_derivative_n_times(self, s, n):
         from sympy import Piecewise, Eq
-        from sympy import Tuple
+        from sympy import Tuple, MatrixExpr
         from sympy.matrices.common import MatrixCommon
-        if isinstance(s, (MatrixCommon, Tuple, Iterable)):
+        if isinstance(s, (MatrixCommon, Tuple, Iterable, MatrixExpr)):
             return super(AtomicExpr, self)._eval_derivative_n_times(s, n)
         if self == s:
             return Piecewise((self, Eq(n, 0)), (1, Eq(n, 1)), (0, True))
diff --git a/sympy/core/function.py b/sympy/core/function.py
@@ -1184,8 +1184,8 @@ def _diff_wrt(self):
     def __new__(cls, expr, *variables, **kwargs):
 
         from sympy.matrices.common import MatrixCommon
-        from sympy import Integer
-        from sympy.tensor.array import Array, NDimArray
+        from sympy import Integer, MatrixExpr
+        from sympy.tensor.array import Array, NDimArray, derive_by_array
         from sympy.utilities.misc import filldedent
 
         expr = sympify(expr)
@@ -1317,6 +1317,9 @@ def __new__(cls, expr, *variables, **kwargs):
                         if not expr.xreplace({v: D}).has(D):
                             zero = True
                             break
+                    elif isinstance(v, MatrixExpr):
+                        zero = False
+                        break
                     elif isinstance(v, Symbol) and v not in free:
                         zero = True
                         break
diff --git a/sympy/core/mul.py b/sympy/core/mul.py
@@ -897,7 +897,7 @@ def _eval_derivative(self, s):
                 # Note: reduce is used in step of Mul as Mul is unable to
                 # handle subtypes and operation priority:
                 terms.append(reduce(lambda x, y: x*y, (args[:i] + [d] + args[i + 1:]), S.One))
-        return reduce(lambda x, y: x+y, terms, S.Zero)
+        return Add.fromiter(terms)
 
     @cacheit
     def _eval_derivative_n_times(self, s, n):
diff --git a/sympy/matrices/expressions/matexpr.py b/sympy/matrices/expressions/matexpr.py
@@ -202,6 +202,23 @@ def _eval_derivative(self, x):
     def _eval_derivative_n_times(self, x, n):
         return Basic._eval_derivative_n_times(self, x, n)
 
+    def _visit_eval_derivative_scalar(self, x):
+        # `x` is a scalar:
+        if x.has(self):
+            return _matrix_derivative(x, self)
+        else:
+            return ZeroMatrix(*self.shape)
+
+    def _visit_eval_derivative_array(self, x):
+        if x.has(self):
+            return _matrix_derivative(x, self)
+        else:
+            from sympy import Derivative
+            return Derivative(x, self)
+
+    def _accept_eval_derivative(self, s):
+        return s._visit_eval_derivative_array(self)
+
     def _entry(self, i, j, **kwargs):
         raise NotImplementedError(
             "Indexing not implemented for %s" % self.__class__.__name__)
@@ -599,24 +616,17 @@ def _postprocessor(expr):
     "Add": [get_postprocessor(Add)],
 }
 
+
 def _matrix_derivative(expr, x):
     from sympy import Derivative
     lines = expr._eval_derivative_matrix_lines(x)
 
-    first = lines[0].first
-    second = lines[0].second
-    higher = lines[0].higher
-
     ranks = [i.rank() for i in lines]
     assert len(set(ranks)) == 1
     rank = ranks[0]
 
     if rank <= 2:
-        return reduce(lambda x, y: x+y, [i.matrix_form() for i in lines])
-        if first != 1:
-            return reduce(lambda x,y: x+y, [lr.first * lr.second.T for lr in lines])
-        elif higher != 1:
-            return reduce(lambda x,y: x+y, [lr.higher for lr in lines])
+        return Add.fromiter([i.matrix_form() for i in lines])
 
     return Derivative(expr, x)
 
@@ -748,8 +758,8 @@ def _eval_derivative_matrix_lines(self, x):
                 transposed=False,
             )]
         else:
-            first=Identity(self.shape[0])
-            second=Identity(self.shape[1])
+            first = Identity(self.shape[0])
+            second = Identity(self.shape[1])
             return [_LeftRightArgs(
                 first=first,
                 second=second,
@@ -848,6 +858,7 @@ def __ne__(self, other):
     def __hash__(self):
         return super(GenericIdentity, self).__hash__()
 
+
 class ZeroMatrix(MatrixExpr):
     """The Matrix Zero 0 - additive identity
 
@@ -902,6 +913,7 @@ def __nonzero__(self):
 
     __bool__ = __nonzero__
 
+
 class GenericZeroMatrix(ZeroMatrix):
     """
     A zero matrix without a specified shape
@@ -937,6 +949,7 @@ def __ne__(self, other):
     def __hash__(self):
         return super(GenericZeroMatrix, self).__hash__()
 
+
 def matrix_symbols(expr):
     return [sym for sym in expr.free_symbols if sym.is_Matrix]
 
@@ -976,6 +989,12 @@ def transpose(self):
     def matrix_form(self):
         if self.first != 1 and self.higher != 1:
             raise ValueError("higher dimensional array cannot be represented")
+        # Remove one-dimensional identity matrices:
+        # (this is needed by `a.diff(a)` where `a` is a vector)
+        if self.first == Identity(1):
+            return self.second.T
+        if self.second == Identity(1):
+            return self.first
         if self.first != 1:
             return self.first*self.second.T
         else:
diff --git a/sympy/matrices/expressions/tests/test_derivatives.py b/sympy/matrices/expressions/tests/test_derivatives.py
@@ -33,7 +33,8 @@ def test_matrix_derivative_non_matrix_result():
 
 def test_matrix_derivative_trivial_cases():
     # Cookbook example 33:
-    assert X.diff(A) == 0
+    # TODO: find a way to represent a four-dimensional zero-array:
+    assert X.diff(A) == Derivative(X, A)
 
 
 def test_matrix_derivative_with_inverse():
@@ -58,6 +59,9 @@ def test_matrix_derivative_with_inverse():
 
 def test_matrix_derivative_vectors_and_scalars():
 
+    assert x.diff(x) == Identity(k)
+    assert x.T.diff(x) == Identity(k)
+
     # Cookbook example 69:
     expr = x.T*a
     assert expr.diff(x) == a
@@ -246,9 +250,23 @@ def test_derivatives_of_complicated_matrix_expr():
 
 def test_mixed_deriv_mixed_expressions():
 
+    expr = 3*Trace(A)
+    assert expr.diff(A) == 3*Identity(k)
+
+    expr = k
+    deriv = expr.diff(A)
+    assert isinstance(deriv, ZeroMatrix)
+    assert deriv == ZeroMatrix(k, k)
+
+    expr = Trace(A)**2
+    assert expr.diff(A) == (2*Trace(A))*Identity(k)
+
     expr = Trace(A)*A
     # TODO: this is not yet supported:
     assert expr.diff(A) == Derivative(expr, A)
 
     expr = Trace(Trace(A)*A)
     assert expr.diff(A) == (2*Trace(A))*Identity(k)
+
+    expr = Trace(Trace(Trace(A)*A)*A)
+    assert expr.diff(A) == (3*Trace(A)**2)*Identity(k)
diff --git a/sympy/matrices/tests/test_matrices.py b/sympy/matrices/tests/test_matrices.py
@@ -2135,7 +2135,8 @@ def test_diff_by_matrix():
     assert A.diff(a) == MutableDenseMatrix([[0, 0], [0, 0]])
 
     B = ImmutableDenseMatrix([a, b])
-    assert A.diff(B) == A.zeros(2)
+    assert A.diff(B) == Array.zeros(2, 1, 2, 2)
+    assert A.diff(A) == Array([[[[1, 0], [0, 0]], [[0, 1], [0, 0]]], [[[0, 0], [1, 0]], [[0, 0], [0, 1]]]])
 
     # Test diff with tuples:
 
