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 evaluates its rhs with `_eval_at(lhs, subdomain=S)`. A `halo=0`
  derivative records S, and multiplies each stencil tap by a branch-free
  integer MIN/MAX indicator of the grid point it reads being in S
  (SubDomain.indicator), which invariant hoisting computes once. Along the
  Dimensions it does not differentiate, it restricts its argument to S at the
  evaluation point. In a sum with such derivatives, the other terms are
  restricted to S at the evaluation point (SubDomain.restrict); factors are
  not.
- The evaluated derivative records its stencil radius (halo_radius), and the
  equation iterates over S grown by the largest one (SubDomain.grow, a
  SubDomain with a parent). Tensor equations aggregate their components.
- 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, including a staggered vector
equation; sums, factors and vector equations; 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_zero_halo(self):
+        return self.halo is not None or self.expr._has_zero_halo
+
     @property
     def is_TimeDependent(self):
         return self.expr.is_TimeDependent
@@ -481,26 +515,21 @@ 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`, which
+        the Derivative records.
         """
-        # No staggering, don't waste time
-        if not self.expr.staggered and not func.staggered:
-            return self
-        # If an x0 already exists or evaluating at the same function (i.e u = u.dx)
-        # do not overwrite it
-        if self.x0 or self.side is not None or func.function is self.expr.function:
-            return self
-        # For basic equation of the form f = Derivative(g, ...) we can just
-        # compare staggering
-        if self.expr.staggered == func.staggered and self.expr.is_Function:
-            return self
-        # Time derivatives are not affected by space staggering
-        if all(d.is_Time for d in self.dims):
-            return self
+        rkw = {}
+        if subdomain is not None and self.halo is not None:
+            rkw['subdomain'] = subdomain
+
+        if not self._is_relocated(func):
+            return self._rebuild(**rkw) if rkw else self
 
         # Check if x0's keys come from a DerivedDimension
         x0 = func.indices_ref.getters
@@ -519,7 +548,7 @@ def _eval_at(self, func, interp_mode='direct', **kwargs):
                     # e.g f.dx(x0={x: x + h_x/2}).subs({x: ix})
                     psubs[sd] = d
                     nx0[sd] = nx0.pop(d)._subs(d, sd)
-        rkw = {'x0': nx0}
+        rkw['x0'] = nx0
         if psubs:
             rkw['subs'] = (psubs,)
 
@@ -534,7 +563,8 @@ def _eval_at(self, func, interp_mode='direct', **kwargs):
                 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)._eval_at(func, interp_mode=interp_mode, **kwargs)
+            args = [self._rebuild(a)._eval_at(func, interp_mode=interp_mode,
+                                              subdomain=subdomain, **kwargs)
                     for a in args]
             return self.expr.func(*args)
         elif self.expr.is_Mul:
@@ -549,6 +579,25 @@ def _eval_at(self, func, interp_mode='direct', **kwargs):
             # the expression as is.
             return self._rebuild(self.expr, **rkw)
 
+    def _is_relocated(self, func):
+        """
+        True if the Derivative must be evaluated at the location of `func`, e.g.
+        with `func` and the argument staggered apart.
+        """
+        # No staggering, don't waste time
+        if not self.expr.staggered and not func.staggered:
+            return False
+        # If an x0 already exists or evaluating at the same function (i.e u = u.dx)
+        # do not overwrite it
+        if self.x0 or self.side is not None or func.function is self.expr.function:
+            return False
+        # For basic equation of the form f = Derivative(g, ...) we can just
+        # compare staggering
+        if self.expr.staggered == func.staggered and self.expr.is_Function:
+            return False
+        # Time derivatives are not affected by space staggering
+        return not all(d.is_Time for d in self.dims)
+
     def _evaluate(self, **kwargs):
         # Evaluate finite-difference.
         # NOTE: `evaluate` and `_eval_fd` split for potential future different
@@ -581,6 +630,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 +654,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,18 +672,26 @@ 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:
             res = res.xreplace(e)
 
+        # With `halo=0`, along the Dimensions it does not differentiate, the
+        # argument is read at the evaluation point, which lies outside the
+        # SubDomain wherever other derivatives extend the equation along them
+        if subdomain is not None:
+            res = subdomain.restrict(res, exclude={d.root for d in self.dims})
+
         return res
 
     def _eval_expand_nest(self, **hints):

devito/finite_differences/differentiable.py

@@ -191,6 +191,24 @@ def _eval_at(self, func, **kwargs):
             for a in self.args  # false positive: lambda is invoked in-place
         ])
 
+    @cached_property
+    def _has_zero_halo(self):
+        """True if the expression has derivatives with `halo=0`."""
+        return any(a._has_zero_halo for a in self._args_diff)
+
+    @cached_property
+    def halo_radius(self):
+        """
+        The largest stencil radius, per root Dimension, of the evaluated
+        derivatives with `halo=0` in the expression. The expression is nonzero up
+        to that many points past the SubDomain their argument is restricted to.
+
+        For example, with `S` a SubDomain restricting `x` and an 8th-order `g`,
+        `g.dx(halo=0)` evaluated in an equation on `S` has a halo radius of
+        `{x: 4}`: the equation must iterate over `S` extended by 4 points.
+        """
+        return merge_halo_radius(a.halo_radius for a in self._args_diff)
+
     def _subs(self, old, new, **hints):
         if old == self:
             return new
@@ -540,6 +558,18 @@ def highest_priority(diff_op, candidates=None):
     return prio_func
 
 
+def merge_halo_radius(radii):
+    """
+    Merge the halo radii `radii`, each a mapping from root Dimension to radius,
+    keeping the largest radius per Dimension.
+    """
+    radius = {}
+    for i in radii:
+        for d, r in i.items():
+            radius[d] = max(radius.get(d, 0), r)
+    return frozendict(radius)
+
+
 class DifferentiableOp(Differentiable):
 
     __sympy_class__ = None
@@ -640,6 +670,19 @@ def __new__(cls, *args, **kwargs):
 
         return super().__new__(cls, *args, **kwargs)
 
+    def _eval_at(self, func, subdomain=None, **kwargs):
+        """
+        Evaluate the sum at the location of `func`.
+
+        The derivatives with `halo=0` extend the sum past `subdomain`, so the
+        other terms are restricted to `subdomain` at the evaluation point.
+        """
+        expr = super()._eval_at(func, subdomain=subdomain, **kwargs)
+        if subdomain is None or not expr.is_Add or not expr._has_zero_halo:
+            return expr
+        halo = {a for a in expr._args_diff if a._has_zero_halo}
+        return self.func(*[a if a in halo else subdomain.restrict(a) for a in expr.args])
+
 
 class Mul(DifferentiableOp, sympy.Mul):
     __sympy_class__ = sympy.Mul
@@ -1276,6 +1319,14 @@ def _subs(self, old, new, **hints):
 
 class DiffDerivative(IndexDerivative, DifferentiableOp):
 
+    __rkwargs__ = IndexDerivative.__rkwargs__ + ('halo_radius',)
+
+    def __new__(cls, *args, halo_radius=None, **kwargs):
+        obj = super().__new__(cls, *args, **kwargs)
+        # With `halo=0`, the stencil radius (see `Differentiable.halo_radius`)
+        obj.halo_radius = frozendict(halo_radius or {})
+        return obj
+
     def _eval_at(self, func, **kwargs):
         # Like EvalDerivative, a DiffDerivative must have already been evaluated
         # at a valid x0 and should not be re-evaluated at a different location
@@ -1291,9 +1342,9 @@ class EvalDerivative(DifferentiableOp, sympy.Add):
 
     is_commutative = True
 
-    __rkwargs__ = ('base',)
+    __rkwargs__ = ('base', 'halo_radius')
 
-    def __new__(cls, *args, base=None, **kwargs):
+    def __new__(cls, *args, base=None, halo_radius=None, **kwargs):
         kwargs['evaluate'] = False
 
         # a+0 -> a
@@ -1309,6 +1360,8 @@ def __new__(cls, *args, base=None, **kwargs):
                 # In some rare cases (rebuild?) base may be obj itself
                 base = base.base
             obj.base = base
+            # With `halo=0`, the stencil radius (see `Differentiable.halo_radius`)
+            obj.halo_radius = frozendict(halo_radius or {})
         except AttributeError:
             # This might happen if e.g. one attempts a (re)construction with
             # one sole argument. The (re)constructed EvalDerivative degenerates

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,34 +227,45 @@ 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`. The derivative then extends past `subdomain` by
+    # the stencil radius
+    halo_radius = {dim.root: indices.radius} if subdomain is not None else {}
+
     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:
+            expr = expr * subdomain.indicator(dim.root, indices.offset(indices.expr))
 
         # Re-evaluate any off-the-grid Functions potentially impacted by the FD
         # unless a pure number
         with suppress(AttributeError):
             expr = expr._evaluate(expand=False)
 
         deriv = DiffDerivative(
-            expr*weights, {dim: indices.free_dim}, deriv_order=deriv_order
+            expr*weights, {dim: indices.free_dim}, deriv_order=deriv_order,
+            halo_radius=halo_radius
         )
     else:
         terms = []
         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))
 
             # Re-evaluate any off-the-grid Functions potentially impacted by the FD
             # unless a pure number
             with suppress(AttributeError):
                 term = term.evaluate
             terms.append(term)
 
-        deriv = EvalDerivative(*terms, base=expr)
+        deriv = EvalDerivative(*terms, base=expr, halo_radius=halo_radius)
 
     return deriv