sympy: Enable add, mul and nest expansion hints

Author: JDBetteridge

devito/finite_differences/derivative.py

@@ -1,11 +1,12 @@
 from collections import OrderedDict
 from collections.abc import Iterable
 from functools import cached_property
+from itertools import chain
 
 import sympy
 
 from .finite_difference import generic_derivative, cross_derivative
-from .differentiable import Differentiable, diffify, interp_for_fd
+from .differentiable import Differentiable, diffify, interp_for_fd, Add
 from .tools import direct, transpose
 from .rsfd import d45
 from devito.tools import (as_mapper, as_tuple, filter_ordered, frozendict, is_integer,
@@ -510,3 +511,33 @@ def _eval_fd(self, expr, **kwargs):
             res = res.xreplace(e)
 
         return res
+
+    def _eval_expand_nest(self, **hints):
+        expr = self.args[0]
+        if isinstance(expr, self.__class__):
+            return self.func(expr.args[0], *[(d, ii)
+                for d, ii in zip(
+                    chain(self.dims, expr.dims),
+                    chain(self.deriv_order, expr.deriv_order)
+                )])
+        else:
+            return self
+
+    def _eval_expand_mul(self, **hints):
+        expr = self.args[0]
+        if isinstance(expr, sympy.Mul):
+            ind, dep = expr.as_independent(*self.dims, as_Mul=True)
+            return ind*self.func(dep, *self.args[1:])
+        else:
+            return self
+
+    def _eval_expand_add(self, **hints):
+        expr = self.args[0]
+        if isinstance(expr, sympy.Add):
+            ind, dep = expr.as_independent(*self.dims, as_Add=True)
+            if isinstance(dep, sympy.Add):
+                return Add(*[self.func(s, *self.args[1:]) for s in dep.args])
+            else:
+                return self.func(dep, *self.args[1:])
+        else:
+            return self