sympy: Add product rule expansion hint and tests

Author: JDBetteridge

devito/finite_differences/derivative.py

@@ -7,7 +7,7 @@
 import sympy
 
 from .finite_difference import generic_derivative, cross_derivative
-from .differentiable import Differentiable, diffify, interp_for_fd, Add
+from .differentiable import Differentiable, diffify, interp_for_fd, Add, Mul
 from .tools import direct, transpose
 from .rsfd import d45
 from devito.tools import (as_mapper, as_tuple, frozendict, is_integer,
@@ -695,4 +695,27 @@ def _eval_expand_product_rule(self, **hints):
         resultant expression for higher derivatives and mixed derivatives is much
         more difficult to implement.
         """
-        raise NotImplementedError('Product rule expansion has not been written')
+        if self.expr.is_Mul and len(self.dims) == 1:
+            if self.deriv_order == (1,):
+                return Add(*[
+                    Mul(*self.expr.args[:ii], self.func(m), *self.expr.args[ii + 1:])
+                    for ii, m in enumerate(self.expr.args)
+                ])
+            elif self.deriv_order == (2,) and len(self.expr.args) == 2:
+                return Add(
+                    Mul(self.func(self.expr.args[0]), self.expr.args[1]),
+                    Mul(
+                        2,
+                        self.func(self.expr.args[0], deriv_order=1),
+                        self.func(self.expr.args[1], deriv_order=1)
+                    ),
+                    Mul(self.expr.args[0], self.func(self.expr.args[1]))
+                )
+            else:
+                # Note: It _is_ possible to implement the product rule for many
+                # more cases, but the number of terms in the resultant expression
+                # will grow.
+                return self
+        else:
+            return self
+

tests/test_derivatives.py

@@ -4,7 +4,7 @@
 
 from devito import (Grid, Function, TimeFunction, Eq, Operator, NODE, cos, sin,
                     ConditionalDimension, left, right, centered, div, grad)
-from devito.finite_differences import Derivative, Differentiable
+from devito.finite_differences import Derivative, Differentiable, diffify
 from devito.finite_differences.differentiable import (Add, EvalDerivative, IndexSum,
                                                       IndexDerivative, Weights,
                                                       DiffDerivative)
@@ -1030,7 +1030,12 @@ class TestExpansion:
     def setup_class(cls):
         cls.grid = Grid(shape=(11,), extent=(1,))
         cls.x = cls.grid.dimensions[0]
-        cls.u = Function(name='u', grid=cls.grid, space_order=4)
+        cls.u = TimeFunction(name='u', grid=cls.grid, space_order=4)
+        cls.t = cls.u.time_dim
+
+        cls.v = Function(name='v', grid=cls.grid, space_order=4)
+        cls.w = Function(name='w', grid=cls.grid, space_order=4)
+        cls.f = Function(name='f', dimensions=(cls.t,), shape=(5,))
 
         # Note that using the `.dx` shortcut method specifies the fd_order kwarg
         a = cls.u.dx
@@ -1140,6 +1145,10 @@ def test_expand_mul(self):
         expanded = Derivative(4*self.u - 5*Derivative(self.u, self.x) + 3, self.x)
         assert self.b.expand() == expanded
 
+        expr = self.u.dx.subs({self.u: 5*self.f*self.u*self.v*self.w})
+        expanded = 5*self.f*Derivative(self.u*self.v*self.w, self.x)
+        assert diffify(expr.expand()) == expanded
+
     def test_expand_add(self):
         """
         Check linearity
@@ -1194,3 +1203,19 @@ def test_nested_orders(self):
         )
         with pytest.raises(ValueError):
             _ = du44.expand(nest=True)
+
+    def test_expand_product_rule(self):
+        """
+        Check the two implemented cases for product rule expansion
+        """
+        expr = self.u.dx.subs({self.u: 5*self.f*self.u*self.v*self.w}, postprocess=False)
+        expanded = 5*self.f*self.v*self.w*Derivative(self.u, self.x) \
+            + 5*self.f*self.v*self.u*Derivative(self.w, self.x) \
+            + 5*self.f*self.w*self.u*Derivative(self.v, self.x)
+        assert diffify(expr.expand(product_rule=True)) == expanded
+
+        expr = self.u.dx2.subs({self.u: 5*self.f*self.u*self.v}, postprocess=False)
+        expanded = 5*self.f*self.v*Derivative(self.u, (self.x, 2)) \
+            + 10*self.f*Derivative(self.v, self.x)*Derivative(self.u, self.x) \
+            + 5*self.f*self.u*Derivative(self.v, (self.x, 2))
+        assert diffify(expr.expand(product_rule=True)) == expanded