dsl: Differentiate a mixed-staggering sum term by term

`Add` reports its first argument's `indices_ref`, so a sum whose terms sit at
different staggered locations names a position only one of them has, and `x0`
gets resolved against it for all of them. The shear strain `v_x.dy + v_y.dx`
of a staggered velocity is the canonical case: both terms land on the cell
corner, so a shift onto it should be a no-op, and instead each picked up a
spurious one.

Differentiation is linear at every order, so split such a sum in
`Derivative._eval_fd`. Relative error on `D(a+b)` against `D(a) + D(b)` was
0.63 at order 0, 1.20 at order 1 and 0.95 at order 2, with `expand=False` at
order 2 returning exactly zero.

`generic_derivative` also short-circuited a zeroth order derivative only when
`x0` was empty, building a stencil around an expression already sitting at
`x0`. `index_at` answers where an expression sits, and both call sites use it.

Author: mloubout

devito/finite_differences/derivative.py

@@ -14,7 +14,7 @@
 from devito.warnings import warn
 
 from .differentiable import Add, Differentiable, Mul, diffify, interp_for_fd
-from .finite_difference import cross_derivative, generic_derivative
+from .finite_difference import cross_derivative, generic_derivative, indices_at
 from .rsfd import d45
 from .tools import direct, transpose
 
@@ -575,6 +575,12 @@ def _eval_fd(self, expr, **kwargs):
              shited derivative.
         - 4: Apply substitutions.
         """
+        # Differentiation is linear, and a sum of terms at different staggered
+        # locations must use it: `Add` reports its first argument's location,
+        # so `x0` would shift the other terms off the point they sat at.
+        if expr.is_Add and any(len(indices_at(expr, d)) > 1 for d in self.dims):
+            return expr.func(*[self._eval_fd(a, **kwargs) for a in expr.args])
+
         # Step 1: Evaluate non-derivative x0. We currently enforce a simple 2nd order
         # interpolation to avoid very expensive finite differences on top of it
         x0_deriv = self._filter_dims(self.x0)

devito/finite_differences/finite_difference.py

@@ -100,6 +100,29 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, side=None, **kw
     return expr
 
 
+def indices_at(expr, dim):
+    """
+    The locations `expr`'s terms sit at along `dim`.
+
+    Terms with no location of their own, a scalar say, contribute none.
+    """
+    indices = set()
+    for i in (expr.args if expr.is_Add else (expr,)):
+        try:
+            indices.add(i.indices_ref[dim])
+        except (AttributeError, KeyError, IndexError, TypeError):
+            continue
+    return indices
+
+
+def index_at(expr, dim):
+    """
+    Where `expr` sits along `dim`, or None if it does not say.
+    """
+    indices = indices_at(expr, dim)
+    return indices.pop() if len(indices) == 1 else None
+
+
 @check_input
 def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
                        coefficients='taylor', expand=True, weights=None, side=None):
@@ -139,8 +162,9 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
     if deriv_order == 1 and fd_order == 2 and side is None:
         fd_order = 1
 
-    # Zeroth order derivative is just the expression itself if not shifted
-    if deriv_order == 0 and not x0:
+    # Zeroth order is the identity when `expr` already sits at `x0`, not a
+    # stencil centred there.
+    if deriv_order == 0 and (not x0 or index_at(expr, dim) == x0.get(dim)):
         return expr
 
     # Enforce stable time coefficients

tests/test_derivatives.py

@@ -1461,3 +1461,43 @@ def test_unevaluated(self):
         assert Derivative(self.x, self.t)
         assert Derivative(self.x, self.y, self.t)
         assert Derivative(self.x, (self.x, 0))
+
+
+@pytest.mark.parametrize('expand', [True, False])
+@pytest.mark.parametrize('deriv_order', [0, 1, 2])
+def test_deriv_sum_mixed_staggering(expand, deriv_order):
+    """
+    A shifted derivative is linear: `D(a + b) == D(a) + D(b)`, at every order.
+
+    Broke for terms at different staggered locations, `Add` reporting only its
+    first argument's.
+    """
+    so = 8
+    grid = Grid(shape=(41, 41), extent=(40., 40.))
+    x, y = grid.dimensions
+
+    vx = Function(name='vx', grid=grid, space_order=so, staggered=x)
+    vy = Function(name='vy', grid=grid, space_order=so, staggered=y)
+    out = Function(name='out', grid=grid, space_order=so, staggered=(x, y))
+
+    rng = np.random.default_rng(3)
+    for f in (vx, vy):
+        f.data[:] = rng.normal(size=f.shape)
+
+    def shifted(expr):
+        return expr.diff(y, deriv_order=deriv_order, fd_order=2,
+                         x0={y: y + y.spacing/2})
+
+    def run(expr):
+        out.data[:] = 0.
+        Operator(Eq(out, expr), opt=('advanced', {'expand': expand})).apply()
+        return np.array(out.data)
+
+    s = slice(so + 3, -(so + 3))
+    together = run(shifted(vx.dy + vy.dx))[s, s]
+    apart = (run(shifted(vx.dy)) + run(shifted(vy.dx)))[s, s]
+
+    assert np.linalg.norm(apart) > 0
+    # float32 reassociation only: the two forms sum the same terms in a
+    # different order
+    assert np.linalg.norm(together - apart) / np.linalg.norm(apart) < 1e-5