api: Add Derivative halo=0 for SubDomain-restricted derivatives and adjoints

`expr.dx(halo=0)` treats `expr` as zero outside the SubDomain of its
equation, and `.T` keeps the flag. Hence `Inc(q_bar, out_bar.dx(halo=0).T,
subdomain=S)` is the adjoint of `Eq(out, q.dx, subdomain=S)`, i.e. D^T R^T with
R the restriction to S, without a user-side zero-padded work field. Equations
without halo=0 derivatives, or without a SubDomain, are evaluated as before.

- An equation with halo=0 derivatives evaluates its rhs with
  `_eval_at(lhs, subdomain=S)`. The halo=0 derivatives record S, and multiply
  each stencil tap by a SubDomainIndicator (SubDomain.indicator) of the grid
  point it reads: 1 inside S, 0 outside, recording the stencil radius
  (IndexSet.radius). The other terms of a sum are restricted to S at the
  evaluation point (SubDomain.restrict); factors are not.
- The equation then iterates over S grown by the largest radius
  (SubDomain.grow, a SubDomain with a parent), and expands the indicators into
  branch-free integer MIN/MAX, which invariant hoisting computes once.
- It works for expanded, unexpanded and staggered derivatives. Derivatives of
  halo=0 derivatives, MultiSubDomains and methods other than FD raise
  NotImplementedError.

Tests: dot tests of `Eq(out, g.dx, subdomain=S)` against
`Inc(g_bar, out_bar.dx(halo=0).T, subdomain=S)` for left/right/middle
SubDomains, first and second derivatives, orders 2/4/8, Eq and Inc, 2D and
unexpanded forms, in 1D and 3D; staggered dot tests against
`-out_bar.dx(halo=0)`, as in elastic adjoints; sums and factors; SubDomain
growth; a forward check; and a one-face CPML operator.

Author: mloubout

devito/finite_differences/derivative.py

@@ -102,7 +102,7 @@ def _fd_priority(self):
 
     __rargs__ = ('expr', '*dims')
     __rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs',
-                   'x0', 'method', 'weights')
+                   'x0', 'method', 'weights', 'halo', 'subdomain')
 
     def __new__(cls, expr, *dims, **kwargs):
         # Validate the input arguments `expr`, `dims` and `deriv_order`
@@ -161,6 +161,8 @@ def __new__(cls, expr, *dims, **kwargs):
         obj._transpose = kwargs.get("transpose", direct)
         obj._method = kwargs.get("method", 'FD')
         obj._weights = cls._process_weights(**kwargs)
+        obj._halo = cls._validate_halo(kwargs.get("halo"))
+        obj._subdomain = kwargs.get("subdomain")
 
         ppsubs = kwargs.get("subs", kwargs.get("_ppsubs", []))
         processed = []
@@ -177,6 +179,16 @@ def __new__(cls, expr, *dims, **kwargs):
 
         return obj
 
+    @staticmethod
+    def _validate_halo(halo):
+        """
+        Validate `halo`. Only None (read the argument everywhere) and 0 (treat
+        the argument as zero outside the equation's SubDomain) are supported.
+        """
+        if halo not in (None, 0):
+            raise ValueError(f"Expected halo=None or halo=0, got halo={halo}")
+        return halo
+
     @staticmethod
     def _validate_expr(expr):
         """
@@ -325,6 +337,8 @@ def _process_weights(cls, **kwargs):
     def __call__(self, x0=None, fd_order=None, side=None, method=None, **kwargs):
         weights = kwargs.get('weights', kwargs.get('w'))
         rkw = {}
+        if 'halo' in kwargs:
+            rkw['halo'] = kwargs['halo']
         if side is not None:
             rkw['side'] = side
         if method is not None:
@@ -457,6 +471,26 @@ def side(self):
     def transpose(self):
         return self._transpose
 
+    @property
+    def halo(self):
+        """
+        None if the argument is read everywhere, 0 if it is treated as zero
+        outside the SubDomain of the equation the Derivative belongs to.
+        """
+        return self._halo
+
+    @property
+    def subdomain(self):
+        """
+        With halo=0, the SubDomain outside of which the argument is treated as
+        zero, set upon evaluation within an equation restricted to it.
+        """
+        return self._subdomain
+
+    @cached_property
+    def _has_halo(self):
+        return self.halo is not None or getattr(self.expr, '_has_halo', False)
+
     @property
     def is_TimeDependent(self):
         return self.expr.is_TimeDependent
@@ -481,12 +515,18 @@ def T(self):
 
         return self._rebuild(transpose=adjoint)
 
-    def _eval_at(self, func, interp_mode='direct', **kwargs):
+    def _eval_at(self, func, interp_mode='direct', subdomain=None, **kwargs):
         """
         Evaluates the derivative at the location of `func`. It is necessary for staggered
         setup where one could have Eq(u(x + h_x/2), v(x).dx)) in which case v(x).dx
         has to be computed at x=x + h_x/2.
+
+        With halo=0, the argument is treated as zero outside `subdomain`.
         """
+        if subdomain is not None and self.halo is not None:
+            return self._rebuild(subdomain=subdomain)._eval_at(
+                func, interp_mode=interp_mode, **kwargs
+            )
         # No staggering, don't waste time
         if not self.expr.staggered and not func.staggered:
             return self
@@ -581,6 +621,13 @@ def _eval_fd(self, expr, **kwargs):
         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])
 
+        # The SubDomain mask read by a halo=0 derivative can't be shifted again
+        if any(d.halo is not None for d in expr.find(Derivative)):
+            raise NotImplementedError(
+                f"{self} differentiates a derivative with halo=0, which is not "
+                "supported"
+            )
+
         # 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)
@@ -598,6 +645,15 @@ def _eval_fd(self, expr, **kwargs):
         # otherwise an IndexSum will returned
         expand = kwargs.get('expand', True)
 
+        # With halo=0, `expr` is treated as zero outside the equation's SubDomain
+        subdomain = self.subdomain
+        if subdomain is not None and (subdomain.is_MultiSubDomain or
+                                      self.method != 'FD'):
+            raise NotImplementedError(
+                f"halo=0 is only supported with method='FD' on a SubDomain, not "
+                f"with method={self.method} on {subdomain}"
+            )
+
         # Step 3: Evaluate FD of the new expression
         if self.method == 'RSFD':
             assert len(self.dims) == 1
@@ -607,13 +663,15 @@ def _eval_fd(self, expr, **kwargs):
             assert self.method == 'FD'
             res = cross_derivative(expr, self.dims, self.fd_order, self.deriv_order,
                                    matvec=self.transpose, x0=x0_deriv, expand=expand,
-                                   side=self.side, weights=self.weights)
+                                   side=self.side, weights=self.weights,
+                                   subdomain=subdomain)
         else:
             assert self.method == 'FD'
             res = generic_derivative(expr, self.dims[0], self.fd_order[0],
                                      self.deriv_order[0], weights=self.weights,
                                      side=self.side, matvec=self.transpose,
-                                     x0=self.x0, expand=expand)
+                                     x0=self.x0, expand=expand,
+                                     subdomain=subdomain)
 
         # Step 4: Apply substitutions
         for e in self._ppsubs:

devito/finite_differences/differentiable.py

@@ -185,11 +185,37 @@ def coefficients(self):
         key = lambda x: coeff_priority.get(x, -1)
         return sorted(coefficients, key=key, reverse=True)[0]
 
-    def _eval_at(self, func, **kwargs):
-        return self.func(*[
-            getattr(a, '_eval_at', lambda x, **kw: a)(func, **kwargs)  # noqa: B023
-            for a in self.args  # false positive: lambda is invoked in-place
-        ])
+    def _eval_at(self, func, subdomain=None, **kwargs):
+        """
+        Evaluate the expression at the location of `func`.
+
+        With a `subdomain`, the derivatives with halo=0 treat their argument as
+        zero outside it, and extend the expression past it. The terms of a sum
+        without such derivatives are then restricted to `subdomain` at the
+        evaluation point.
+        """
+        args = []
+        for a in self.args:
+            try:
+                eval_at = a._eval_at
+            except AttributeError:
+                args.append(a)
+                continue
+            if subdomain is None:
+                args.append(eval_at(func, **kwargs))
+            elif getattr(a, '_has_halo', False):
+                args.append(eval_at(func, subdomain=subdomain, **kwargs))
+            elif self.is_Add:
+                args.append(subdomain.restrict(eval_at(func, **kwargs)))
+            else:
+                # A factor, e.g. `c` in `c*f.dx(halo=0)`, is not restricted
+                args.append(eval_at(func, **kwargs))
+        return self.func(*args)
+
+    @cached_property
+    def _has_halo(self):
+        """True if the expression has derivatives with halo=0."""
+        return any(getattr(a, '_has_halo', False) for a in self.args)
 
     def _subs(self, old, new, **hints):
         if old == self:

devito/finite_differences/finite_difference.py

@@ -125,7 +125,8 @@ def index_at(expr, dim):
 
 @check_input
 def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
-                       coefficients='taylor', expand=True, weights=None, side=None):
+                       coefficients='taylor', expand=True, weights=None, side=None,
+                       subdomain=None):
     """
     Arbitrary-order derivative of a given expression.
 
@@ -151,6 +152,10 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
     expand : bool, optional, default=True
         If True, the derivative is fully expanded as a sum of products,
         otherwise an IndexSum is returned.
+    subdomain : SubDomain, optional
+        If given, `expr` is treated as zero outside `subdomain` (see
+        `Derivative`'s `halo`): each stencil tap is multiplied by the
+        indicator of `subdomain` at the point it reads.
 
     Returns
     -------
@@ -171,7 +176,7 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
     coefficients = 'taylor' if dim.is_Time else expr.coefficients
 
     return make_derivative(expr, dim, fd_order, deriv_order, side,
-                           matvec, x0, coefficients, expand, weights)
+                           matvec, x0, coefficients, expand, weights, subdomain)
 
 
 # Backward compatibility
@@ -180,7 +185,7 @@ def first_derivative(expr, dim, fd_order, **kwargs):
 
 
 def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coefficients,
-                    expand, weights=None):
+                    expand, weights=None, subdomain=None):
     # Always expand time derivatives to avoid issue with buffering and streaming.
     # Time derivative are almost always short stencils and won't benefit from
     # unexpansion in the rare case the derivative is not evaluated for time stepping.
@@ -222,13 +227,20 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
     if callable(expand):
         expand = expand(dim)
 
+    # 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`
+
     if not expand and indices.expr is not None:
         weights = Weights(name='w', dimensions=indices.free_dim,
                           initvalue=weights, dtype=expr.dtype)
 
         # Inject the StencilDimension
         # E.g. `x + i*h_x` into `f(x)` s.t. `f(x + i*h_x)`
         expr = expr.shift(dim, indices.expr - dim)
+        if subdomain is not None:
+            offset = indices.offset(indices.expr)
+            expr = expr * subdomain.indicator(dim.root, offset, indices.radius)
 
         # Re-evaluate any off-the-grid Functions potentially impacted by the FD
         # unless a pure number
@@ -243,6 +255,9 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
         for i, c in zip(indices, weights, strict=True):
             # The FD term
             term = expr.shift(dim, i - dim) * c
+            if subdomain is not None:
+                term = term * subdomain.indicator(dim.root, indices.offset(i),
+                                                  indices.radius)
 
             # Re-evaluate any off-the-grid Functions potentially impacted by the FD
             # unless a pure number

devito/finite_differences/tools.py

@@ -1,4 +1,4 @@
-from functools import partial, wraps
+from functools import cached_property, partial, wraps
 from itertools import product
 
 import numpy as np
@@ -177,6 +177,18 @@ def __repr__(self):
     def spacing(self):
         return self.dim.spacing
 
+    def offset(self, index):
+        """
+        Offset, in grid points, of the point `index` from `dim`: e.g. 2 for
+        `x + 2*h_x`, or `i0` for `x + i0*h_x` in unexpanded form.
+        """
+        return sympify((index - self.dim) / self.spacing)
+
+    @cached_property
+    def radius(self):
+        """Largest distance, in grid points, between `dim` and the points."""
+        return max(abs(self.offset(i)) for i in self)
+
     def transpose(self):
         """
         Transpose the IndexSet.

devito/types/equation.py

@@ -125,18 +125,48 @@ def _evaluate(self, **kwargs):
         """
         if self._interp_mode is not None:
             kwargs['interp_mode'] = self._interp_mode
+
+        subdomain = self.subdomain
+        restrict = subdomain is not None and getattr(self.rhs, '_has_halo', False)
+
+        # Derivatives with halo=0 treat their argument as zero outside the
+        # SubDomain, and extend the rhs past it
+        at = {'subdomain': subdomain} if restrict else {}
         try:
             lhs = self.lhs._evaluate(**kwargs)
-            rhs = self.rhs._eval_at(self.lhs, **kwargs)._evaluate(**kwargs)
+            rhs = self.rhs._eval_at(self.lhs, **at, **kwargs)._evaluate(**kwargs)
         except AttributeError:
             lhs, rhs = self._evaluate_args(**kwargs)
-        eq = self.func(lhs, rhs, subdomain=self.subdomain,
+
+        if restrict:
+            rhs, subdomain = self._extend(rhs, subdomain)
+
+        eq = self.func(lhs, rhs, subdomain=subdomain,
                        coefficients=self.substitutions,
                        implicit_dims=self._implicit_dims,
                        interp_mode=self._interp_mode)
 
         return eq
 
+    @staticmethod
+    def _extend(rhs, subdomain):
+        """
+        The evaluated `rhs`, with its SubDomainIndicators expanded, and the region
+        over which it is evaluated: `subdomain` grown by the largest radius of the
+        stencils reading them.
+        """
+        from devito.symbolics import search  # noqa
+
+        indicators = search(rhs, lambda e: getattr(e, 'is_SubDomainIndicator', False))
+
+        radius = {}
+        for i in indicators:
+            if i.radius > 0:
+                radius[i.dim] = max(radius.get(i.dim, 0), i.radius)
+        region = subdomain.grow(radius) if radius else subdomain
+
+        return rhs.xreplace({i: i.branchless for i in indicators}), region
+
     @property
     def _flatten(self):
         """