compiler: Patch ideriv._subs

Author: FabioLuporini

devito/finite_differences/differentiable.py

@@ -948,6 +948,20 @@ def _evaluate(self, **kwargs):
 
         return EvalDerivative(*expr.args, base=self.base)
 
+    def _subs(self, old, new, **hints):
+        # We have to work around SymPy's weak implementation of `subs` when
+        # it gets to replacing sub-operations such as `a*b*c` (i.e., potentially
+        # `self`'s `base`) within say `a*b*c*w[i0]` (i.e., the corresponding
+        # `self.expr`), because depending on the complexity of `a/b/c`, SymPy
+        # may fail to identify the sub-expression to be replaced (note: if
+        # `a/b/c` are atoms or Indexeds, it's generally fine)
+
+        if not old.is_Mul or \
+           old is not self.base:
+            return super()._subs(old, new, **hints)
+
+        return self._rebuild(new * self.weights)
+
 
 class DiffDerivative(IndexDerivative, DifferentiableOp):
     pass

tests/test_symbolics.py

@@ -717,6 +717,25 @@ def test_canonical_ordering_of_weights():
     assert ccode(cf*wi) == 'f[x][y + i0][z].x*w0[i0]'
 
 
+def test_ideriv_subs_complex():
+    grid = Grid(shape=(3, 3))
+    x, _ = grid.dimensions
+
+    f = Function(name='f', grid=grid, space_order=4)
+    g = f.func(name='g')
+    h = f.func(name='h')
+    b = f.func(name='b')
+
+    ideriv = (f*g*h).dx._evaluate(expand=False)
+
+    i0, = ideriv.dimensions
+    base = ideriv.base
+
+    v = ideriv._subs(base, b._subs(x, x + i0))
+
+    assert str(v) == 'DiffDerivative(w(i0)*b(x + i0, y), (i0))'
+
+
 def test_symbolic_printing():
     b = Symbol('b')