dsl: Fix up functions on subdomains staggered derivs and cross derivs with halo=0

devito/finite_differences/finite_difference.py

@@ -6,7 +6,7 @@
 
 from devito.logger import warning
 
-from .differentiable import DiffDerivative, EvalDerivative, Weights
+from .differentiable import DiffDerivative, EvalDerivative, Weights, merge_growth
 from .tools import (
     centered, check_input, direct, fd_weights_registry, generate_indices, left,
     process_weights, right, transpose
@@ -236,8 +236,12 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
     # With a `subdomain`, `expr` is treated as zero outside it: every stencil tap
     # is multiplied by the indicator of the point it reads, 1 if `subdomain`
     # spans the whole of `dim`. The derivative then extends past `subdomain` by
-    # the stencil radius
-    growth = {dim.root: indices.radius} if subdomain is not None else {}
+    # the stencil radius, on top of the growth of `expr` itself (e.g., the inner
+    # derivative of a cross derivative)
+    if subdomain is not None:
+        growth = merge_growth([expr._growth, {dim.root: indices.radius}])
+    else:
+        growth = {}
 
     if not expand and indices.expr is not None:
         weights = Weights(name='w', dimensions=indices.free_dim,

devito/finite_differences/tools.py

@@ -299,6 +299,10 @@ def generate_indices(expr, dim, order, side=None, matvec=None, x0=None, nweights
         iexpr = x0 + d * dim.spacing
         return IndexSet(dim, expr=iexpr), x0
 
+    # Normalize x0 `Dimension` to match index if the latter is a `SubDimension` of
+    # the former
+    x0 = x0.subs({d.root: d for d in expr.dimensions if d.root is dim and d.is_Sub})
+
     # Evaluation point relative to the expression's grid
     mid = (x0 - expr.indices_ref[dim]).subs({dim: 0, dim.spacing: 1})
 

devito/types/dense.py

@@ -1165,16 +1165,18 @@ def __staggered_setup__(cls, dimensions, staggered=None, **kwargs):
             processed = staggered
         else:
             # Staggering is not NODE or CELL or None
-            # therefore it's a tuple of dimensions
+            # therefore it's a tuple of dimensions. These may be the roots of
+            # `dimensions` (e.g., `staggered=x` for a Function on a SubDomain)
+            staggered = set(as_tuple(staggered))
             processed = []
             for d in dimensions:
-                if d in as_tuple(staggered):
+                if {d, d.root} & staggered:
                     processed.append(sympy.S.One)
-                elif -d in as_tuple(staggered):
+                elif {-d, -d.root} & staggered:
                     processed.append(sympy.S.NegativeOne)
                 else:
                     processed.append(sympy.S.Zero)
-        return Staggering(*processed, getters=dimensions)
+        return Staggering(*processed, getters=[d.root for d in dimensions])
 
     @classmethod
     def __indices_setup__(cls, *args, **kwargs):

tests/test_derivatives.py

@@ -1861,7 +1861,7 @@ def test_adjoint_unexpanded_3d(self):
 
         assert self.dot_test(grid, subdomain, 'dy', 8, opt=opt) < 1e-12
 
-    @pytest.mark.parametrize('deriv', ['dx', 'dy'])
+    @pytest.mark.parametrize('deriv', ['dx', 'dy', 'dxdy'])
     def test_adjoint_2d(self, deriv):
         grid = Grid(shape=(20, 22), extent=(19., 21.), dtype=np.float64)
         subdomain = Region('s', {'x': ('left', 6), 'y': ('middle', 4, 5)}, grid=grid)

tests/test_subdomains.py

@@ -6,8 +6,8 @@
 
 from conftest import assert_structure, opts_tiling
 from devito import (
-    Border, Buffer, ConditionalDimension, Constant, Eq, Function, Grid, Lt, Operator,
-    SparseFunction, SparseTimeFunction, SubDimension, SubDomain, SubDomainSet,
+    Border, Buffer, ConditionalDimension, Constant, Derivative, Eq, Function, Grid, Lt,
+    Operator, SparseFunction, SparseTimeFunction, SubDimension, SubDomain, SubDomainSet,
     TensorFunction, TimeFunction, VectorFunction, solve
 )
 from devito.exceptions import InvalidArgument
@@ -1601,6 +1601,40 @@ def define(self, dimensions):
         assert np.all(np.isclose(fdx.data[:], gdx.data[2:-2, 3:-1]))
         assert np.all(np.isclose(fdy.data[:], gdy.data[2:-2, 3:-1]))
 
+    @pytest.mark.parametrize('s_o', [2, 4])
+    @pytest.mark.parametrize('axis', [0, 1])
+    def test_derivatives_staggered(self, s_o, axis):
+        """
+        Derivatives of Functions on SubDomains can be evaluated at points staggered
+        along the derivative Dimension
+        """
+
+        class Middle(SubDomain):
+
+            name = 'middle'
+
+            def define(self, dimensions):
+                x, y = dimensions
+                return {x: ('middle', 2, 2), y: ('middle', 3, 1)}
+
+        grid = Grid(shape=(10, 10), extent=(9., 9.))
+        d = grid.dimensions[axis]
+        mid = Middle(grid=grid)
+
+        f = Function(name='f', grid=mid, space_order=s_o)
+        g = Function(name='g', grid=grid, space_order=s_o)
+        fd, gd = (Function(name=i, grid=grid, staggered=d) for i in ('fd', 'gd'))
+
+        f.data[:] = np.random.default_rng(0).random(f.shape)
+        g.data[2:-2, 3:-1] = f.data
+
+        Operator([Eq(fd, Derivative(f, d), subdomain=mid),
+                  Eq(gd, Derivative(g, d), subdomain=mid)])()
+
+        # Same stencil, so bitwise identical
+        assert np.any(fd.data != 0)
+        assert np.all(fd.data == gd.data)
+
     def test_operator_shortcuts(self):
         """
         Test that shortcuts for Laplacian, divergence, etc all work with Functions
@@ -1982,14 +2016,14 @@ def define(self, dimensions):
                 return {x: ('left', 2), y: y, z: z}
 
         sd1 = SD1(grid=grid)
-        f = TimeFunction(name='f', grid=grid, time_order=1, space_order=2,
-                         staggered=(x,))
+        f = TimeFunction(name='f', grid=grid, time_order=1, space_order=2)
         g = TimeFunction(name='g', grid=sd1, time_order=1, space_order=2,
                          staggered=(x,))
 
         eq = Eq(g, g + f.dx)
         eqe = eq.evaluate
-        assert eqe.rhs == g + f.dx(x0=x).evaluate._subs(x, g.dimensions[1])
+        x0 = x + x.spacing/2
+        assert eqe.rhs == g + f.dx(x0=x0).evaluate._subs(x, g.dimensions[1])
 
 
 class TestSubDomainArguments: