api: fix staggered evaluation for add/mul

Author: mloubout

devito/finite_differences/derivative.py

@@ -530,9 +530,12 @@ def _eval_at(self, func):
             # it into `u(x + h_x/2).dx` and `v(x).dx`, since they require
             # different FD indices
             mapper = as_mapper(self.expr._args_diff, lambda i: i.staggered)
+            if len(mapper) == 1:
+                # All terms have the same staggering, we can use expr as is
+                return self._rebuild(self.expr, **rkw)
             args = [self.expr.func(*v) for v in mapper.values()]
             args.extend([a for a in self.expr.args if a not in self.expr._args_diff])
-            args = [self._rebuild(a, **rkw) for a in args]
+            args = [self._rebuild(a)._eval_at(func) for a in args]
             return self.expr.func(*args)
         elif self.expr.is_Mul:
             # For Mul, We treat the basic case `u(x + h_x/2) * v(x) which is what appear

devito/finite_differences/differentiable.py

@@ -1134,6 +1134,13 @@ def _(expr, x0, **kwargs):
     from devito.finite_differences.derivative import Derivative
     x0_expr = {d: v for d, v in x0.items() if v is not expr.indices_ref[d]}
     if expr.is_parameter:
+        # Parameter might not have been evaluated at x0 yet
+        # E.g., in expressions such as u[x]*f[x] evaluated at x+dx/2
+        # `f` will not be evaluated at it since gather_for_diff will pick `u` as
+        # higher priority function.
+        for d, v in x0.items():
+            if expr.indices[d] is d:
+                expr = expr._subs(d, v)
         return expr
     elif x0_expr:
         dims = tuple((d, 0) for d in x0_expr)

tests/test_derivatives.py

@@ -780,6 +780,44 @@ def test_param_stagg_inner(self):
         eqne = eqn.evaluate.rhs
         assert simplify(eqne - (p._subs(y, yp).evaluate * f).dx(x0=xp).evaluate) == 0
 
+    def test_param_stagg_add(self):
+        space_order = 2
+        nx, ny = 5, 5
+        extent = (nx-1, ny-1)
+
+        grid = Grid(shape=(nx, ny), extent=extent)
+        x, y = grid.dimensions
+        yp = y + y.spacing / 2
+        xp = x + x.spacing / 2
+
+        x, y = grid.dimensions
+
+        vx = TimeFunction(name="vx", grid=grid, space_order=space_order,
+                          time_order1=1, staggered=x)
+        txx = TimeFunction(name="txx", grid=grid, space_order=space_order,
+                           time_order=1, staggered=NODE)
+        txy = TimeFunction(name="txy", grid=grid, space_order=space_order,
+                           time_order=1, staggered=(x, y))
+        c11 = Function(name="c11", grid=grid, space_order=space_order, parameter=True)
+        c66 = Function(name="c66", grid=grid, space_order=space_order, parameter=True)
+
+        eq0 = Eq(vx, (c66 * txy).dy)
+        eq1 = Eq(vx, (c11 * txx).dy)
+        eq2 = Eq(vx, (c11 * txx + c66 * txy).dy)
+
+        # C66 is a paramater. Expects to evaluate c66 at xp then the derivative at yp
+        # and the derivative will interpolate txy at xp
+        expect0 = (c66.subs({x: xp, y: yp}).evaluate * txy).dy.evaluate
+        assert simplify(eq0.evaluate.rhs - expect0) == 0
+
+        # C11 is a paramater and txy is staggered in x.
+        # Expects to evaluate c11 and txy xp then the derivative at yp
+        expect1 = (c11._subs(x, xp).evaluate * txx._subs(x, xp).evaluate).dy.evaluate
+        assert simplify(eq1.evaluate.rhs - expect1) == 0
+
+        # Addition should apply the same logic as above for each term
+        assert simplify(eq2.evaluate.rhs - (expect1 + expect0)) == 0
+
 
 class TestTwoStageEvaluation:
 

tests/test_staggered_utils.py

@@ -99,7 +99,7 @@ def test_gather_for_diff(expr, expected):
 
 @pytest.mark.parametrize('expr, expected', [
     ('((a + b).dx._eval_at(a)).is_Add', 'True'),
-    ('(a + b).dx._eval_at(a)', 'a.dx(x0=a.indices_ref.getters) + b.dx._eval_at(a)'),
+    ('(a + b).dx._eval_at(a)', 'a.dx + b.dx._eval_at(a)'),
     ('(a*b).dx._eval_at(a).expr', 'a.subs({x: x0}) * b'),
     ('(a * b.dx).dx._eval_at(b).expr._eval_deriv ',
      'a.subs({x: x0}) * b.dx.evaluate')])