Merge pull request #3001 from devitocodes/symmetric-interp-gradients

api: Symmetric interp gradients

Author: mloubout

devito/finite_differences/differentiable.py

@@ -15,7 +15,9 @@
     # Moved in 1.13
     from sympy.core.basic import ordering_of_classes
 
-from devito.finite_differences.interpolation import interp_at, post_x0_indices
+from devito.finite_differences.interpolation import (
+    interp_at, interp_mapper, post_x0_indices
+)
 from devito.finite_differences.tools import coeff_priority, make_shift_x0
 from devito.logger import warning
 from devito.tools import (
@@ -487,16 +489,23 @@ def has_free(self, *patterns):
             return all(i in self.free_symbols for i in patterns)
 
 
-def highest_priority(diff_op):
-    if not diff_op._args_diff:
+def highest_priority(diff_op, candidates=None):
+    """
+    The Function whose location a product should be evaluated at.
+
+    `candidates` restricts the choice to a subset of the operands; without it
+    the whole expression's differentiable arguments are considered.
+    """
+    args_diff = diff_op._args_diff if candidates is None else tuple(candidates)
+    if not args_diff:
         return diff_op
 
     # We want to get the object with highest priority
     # We also need to make sure that the object with the largest
     # set of dimensions is used when multiple ones with the same
     # priority appear
     prio = lambda x: (getattr(x, '_fd_priority', 0), len(x.dimensions))
-    prio_func = sorted(diff_op._args_diff, key=prio, reverse=True)[0]
+    prio_func = sorted(args_diff, key=prio, reverse=True)[0]
 
     # The highest priority must be a Function
     if not isinstance(prio_func, AbstractFunction):
@@ -664,6 +673,27 @@ def _gather_for_diff(self):
             other = self.func(*other)._eval_at(highest_priority(self))
             return self.func(other, *derivs)
 
+    @classmethod
+    def _off_func(cls, a, func):
+        """
+        Whether evaluating `a` at `func`'s location takes an interpolation.
+
+        A derivative is where its `x0` puts it, not where the field it
+        differentiates lives: `div(v)` of a staggered velocity lands on the
+        node, and reading its operand's staggering instead would re-associate a
+        product that is already node-centred and interpolate it for nothing.
+        A composite is off only if one of its own operands is, which is what
+        makes a sum of derivatives -- a divergence, a trace -- read as the
+        node-centred quantity it is.
+        """
+        if isinstance(a, sympy.Derivative):
+            source = post_x0_indices(a, func)
+        elif isinstance(a, AbstractFunction) or not getattr(a, '_args_diff', ()):
+            source = a.indices_ref
+        else:
+            return any(cls._off_func(i, func) for i in a._args_diff)
+        return bool(interp_mapper(source, func.indices_ref, a.dimensions))
+
     def _eval_at(self, func, interp_mode='direct', **kwargs):
         """
         Evaluate a Mul at the location of `func`.
@@ -674,36 +704,52 @@ def _eval_at(self, func, interp_mode='direct', **kwargs):
           independently evaluated at `func`'s location via
           `Differentiable._eval_at`.
 
-        - `interp_mode='symmetric'`: when every Differentiable factor has a
-          staggering different from `func`'s, apply the `I * (a * I^T * b)`
-          form:
-
-            1. Pick a `block` location -- the highest-priority factor's
-               staggering (NODE is the highest priority, so coefficient-like
-               NODE factors win, as in the `I * C * I^T` elastic stiffness
-               pattern). Each factor not at the block is brought there via
-               `I^T` (an explicit 0-order FD interpolation operator).
-               Derivatives additionally set `x0` on their own derivative
-               dimensions to `func`'s indices.
+        - `interp_mode='symmetric'`: the product is formed *away* from `func`
+          and closed with a single interpolation, the `I * (a * I^T * b)` form:
+
+            1. Pick a `block` location -- the highest-priority staggering
+               among the factors that are not already at `func`'s (NODE is the
+               highest priority, so coefficient-like NODE factors win, as in
+               the `I * C * I^T` elastic stiffness pattern). Each factor not at
+               the block is brought there via `I^T` (an explicit 0-order FD
+               interpolation operator). Derivatives additionally set `x0` on
+               their own derivative dimensions to `func`'s indices.
             2. The product is formed at `block`'s location.
             3. The whole product is interpolated to `func` via `I` (an
                explicit 0-order FD operator).
 
-          When the trigger does not hold (e.g. some factor already matches
-          `func`'s staggering), we fall back to `direct`.
+          It takes *two* factors away from `func` to have something to
+          re-associate. With one, the single interpolation `direct` puts on it
+          is already the transpose-consistent form -- it is what makes the
+          `i, j` entry of a stiffness matrix the transpose of its `j, i` entry
+          -- so we fall back. With two, `direct` would interpolate each of them
+          separately, and `I(a)*I(b)` is not `I(a*b)`: the discretized operator
+          stops being the transpose of itself, which is invisible in a forward
+          simulation and shows up as a first-order gradient in an adjoint one.
+
+          Which factors count is the other half of it. A derivative is where
+          its `x0` puts it, not where its operand lives, so a product around a
+          `div(v)` of a staggered velocity -- node-centred however its operands
+          are staggered -- is left alone rather than re-associated onto the
+          operands' location, which would replace a compact stencil with an
+          interpolated one twice as wide.
         """
         if interp_mode != 'symmetric':
             return super()._eval_at(func, **kwargs)
 
         diff, other = split(self.args, lambda a: isinstance(a, Differentiable))
 
-        # Symmetric form requires every Differentiable factor to differ from
-        # func; otherwise direct evaluation is cleaner and equivalent.
-        if len(diff) < 2 or \
-           any(a.staggered == func.staggered for a in diff):
-            return super()._eval_at(func, **kwargs)
+        # A single factor cannot be re-associated, and with everything already
+        # on `func` there is no interpolation to place.  The mode still has to
+        # travel down: a product that cannot itself be re-associated is
+        # routinely wrapped around one that can (a `dt` scaling, a sum of
+        # per-component contractions), and dropping the mode here would silently
+        # evaluate all of that in `direct`.
+        off_func = [a for a in diff if self._off_func(a, func)]
+        if len(off_func) < 2:
+            return super()._eval_at(func, interp_mode=interp_mode, **kwargs)
 
-        block_indices = highest_priority(self).indices_ref
+        block_indices = highest_priority(self, candidates=off_func).indices_ref
 
         # Bring each factor to block's location (I^T where needed)
         new_factors = list(other)
@@ -1119,9 +1165,11 @@ def __new__(cls, *args, base=None, **kwargs):
         except AttributeError:
             # This might happen if e.g. one attempts a (re)construction with
             # one sole argument. The (re)constructed EvalDerivative degenerates
-            # to an object of different type, in classic SymPy style. That's fine
+            # to an object of different type, in classic SymPy style. That's
+            # fine -- and a single argument that is itself a sum is the same
+            # story: a zero-order derivative whose weights collapse to one is
+            # the identity, so it comes back as the sum it was applied to.
             assert len(args) <= 1
-            assert not obj.is_Add
             return obj
 
         return obj

devito/types/basic.py

@@ -1,7 +1,7 @@
 import abc
 import inspect
 import warnings
-from contextlib import suppress
+from contextlib import contextmanager, suppress
 from ctypes import POINTER, Structure, _Pointer, c_char, c_char_p
 from functools import cached_property, reduce
 from operator import mul
@@ -1471,6 +1471,22 @@ def __getnewargs_ex__(self):
         return args, kwargs
 
 
+@contextmanager
+def ignore_non_expr_deprecation():
+    """
+    Suppress sympy's deprecation of non-`Expr` entries in a Matrix.
+
+    A Devito tensor is a Matrix of Devito objects, some of which are legitimately
+    not `Expr` (for example a serialization string), so the deprecation does not
+    apply to us. Sympy emits it from `_dod_to_DomainMatrix`, reached whenever a
+    tensor is built from components, and from `_unify_element_sympy`, reached on
+    scalar multiplication, so both are wrapped in this.
+    """
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
+        yield
+
+
 class AbstractTensor(sympy.ImmutableDenseMatrix, Basic, Pickable, Evaluable):
 
     """
@@ -1563,23 +1579,29 @@ def __subfunc_setup__(cls, *args, **kwargs):
         return []
 
     @classmethod
-    def _sympify(self, arg):
+    def _sympify(cls, arg):
         # This is used internally by sympy to process arguments at rebuilt. And since
-        # some of our properties are non-sympyfiable we need to have a fallback
+        # some of our properties are non-sympyfiable we need to have a fallback.
+        # `strict` so that strings are left alone rather than parsed into Symbols,
+        # while plain numbers are turned into `Expr` as sympy expects (a Matrix
+        # holding non-`Expr` entries, such as a plain `int` 0, is deprecated)
         try:
-            # Pure sympy object
-            return arg._sympy_()
-        except AttributeError:
+            return sympy.sympify(arg, strict=True)
+        except sympy.SympifyError:
             return arg
 
     @classmethod
-    def _eval_from_dok(cls, rows, cols, dok):
-        with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore",
-                category=SymPyDeprecationWarning
-            )
-            return super()._eval_from_dok(rows, cols, dok)
+    def _dod_to_DomainMatrix(cls, rows, cols, dod, types):
+        # Entry point of every matrix construction, from either a flat list
+        # (`_new`) or a dok (`_eval_from_dok`)
+        with ignore_non_expr_deprecation():
+            return super()._dod_to_DomainMatrix(rows, cols, dod, types)
+
+    @classmethod
+    def _unify_element_sympy(cls, rep, element):
+        # Entry point of scalar multiplication, e.g. `lam * tau`
+        with ignore_non_expr_deprecation():
+            return super()._unify_element_sympy(rep, element)
 
     @property
     def grid(self):

devito/types/equation.py

@@ -36,6 +36,15 @@ class Eq(sympy.Eq, Evaluable, Pickable):
         An ordered list of Dimensions that do not explicitly appear in either the
         left-hand side or in the right-hand side, but that should be honored when
         constructing an Operator.
+    interp_mode : str, optional, default=None
+        Overrides the Operator's `sym_opt={'interp-mode': ...}` for this
+        equation only. An Operator generally wants one mode -- `'direct'`, which
+        keeps finite-difference stencils compact -- while a single equation in
+        it may need `'symmetric'`, which re-associates a product of staggered
+        operands so that the discretized operator is the transpose of itself.
+        A gradient accumulated alongside the propagation it is the adjoint of is
+        the typical case: the stencils must stay compact, but the accumulation
+        has to be an exact transpose or the gradient degrades to first order.
 
     Examples
     --------
@@ -60,10 +69,10 @@ class Eq(sympy.Eq, Evaluable, Pickable):
     is_Reduction = False
 
     __rargs__ = ('lhs', 'rhs')
-    __rkwargs__ = ('subdomain', 'coefficients', 'implicit_dims')
+    __rkwargs__ = ('subdomain', 'coefficients', 'implicit_dims', 'interp_mode')
 
     def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None,
-                implicit_dims=None, **kwargs):
+                implicit_dims=None, interp_mode=None, **kwargs):
         if coefficients is not None:
             _ = deprecations.coeff_warn
         kwargs['evaluate'] = False
@@ -76,9 +85,15 @@ def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None,
         obj._subdomain = subdomain
         obj._substitutions = coefficients
         obj._implicit_dims = as_tuple(implicit_dims)
+        obj._interp_mode = interp_mode
 
         return obj
 
+    @property
+    def interp_mode(self):
+        """Per-equation override of the Operator's `interp-mode`, or None."""
+        return self._interp_mode
+
     @classmethod
     def _apply_coeffs(cls, expr, coefficients):
         """
@@ -108,14 +123,17 @@ def _evaluate(self, **kwargs):
 
         The RHS of the Equation is evaluated at the indices of the LHS if required.
         """
+        if self._interp_mode is not None:
+            kwargs['interp_mode'] = self._interp_mode
         try:
             lhs = self.lhs._evaluate(**kwargs)
             rhs = self.rhs._eval_at(self.lhs, **kwargs)._evaluate(**kwargs)
         except AttributeError:
             lhs, rhs = self._evaluate_args(**kwargs)
         eq = self.func(lhs, rhs, subdomain=self.subdomain,
                        coefficients=self.substitutions,
-                       implicit_dims=self._implicit_dims)
+                       implicit_dims=self._implicit_dims,
+                       interp_mode=self._interp_mode)
 
         return eq