api: fix symmetric interp mode and add lots of tests

mloubout · mloubout · commit be046cb3bd2d · 2026-05-11T23:23:55.000-04:00
diff --git a/devito/core/cpu.py b/devito/core/cpu.py
@@ -87,7 +87,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Code generation options for derivatives
         o['expand'] = oo.pop('expand', cls.EXPAND)
-        o['eval-mul-first'] = oo.pop('eval-mul-first', cls.MUL_FIRST)
+        o['interp-mode'] = oo.pop('interp-mode', cls.INTERP_MODE)
         o['deriv-collect'] = oo.pop('deriv-collect', cls.DERIV_COLLECT)
         o['deriv-schedule'] = oo.pop('deriv-schedule', cls.DERIV_SCHEDULE)
         o['deriv-unroll'] = oo.pop('deriv-unroll', False)
diff --git a/devito/core/gpu.py b/devito/core/gpu.py
@@ -102,7 +102,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Code generation options for derivatives
         o['expand'] = oo.pop('expand', cls.EXPAND)
-        o['eval-mul-first'] = oo.pop('eval-mul-first', cls.MUL_FIRST)
+        o['interp-mode'] = oo.pop('interp-mode', cls.INTERP_MODE)
         o['deriv-collect'] = oo.pop('deriv-collect', cls.DERIV_COLLECT)
         o['deriv-schedule'] = oo.pop('deriv-schedule', cls.DERIV_SCHEDULE)
         o['deriv-unroll'] = oo.pop('deriv-unroll', False)
diff --git a/devito/core/operator.py b/devito/core/operator.py
@@ -125,10 +125,24 @@ class BasicOperator(Operator):
     finite-difference derivatives.
     """
 
-    MUL_FIRST = False
+    INTERP_MODE = 'direct'
     """
-    When evaluating expressions location, prioritize multiplication
-    operations.
+    Interpolation mode used by `Mul._eval_at` when projecting a multi-factor
+    expression onto a target staggered location:
+
+    * `'direct'` (default): each factor is shifted to `func`'s location
+      independently (`Function._eval_at` per arg). Cheapest stencil; the
+      mode to pick unless you need an explicitly self-adjoint discretization.
+
+    * `'symmetric'`: when every factor lives at a staggering different from
+      `func`'s, the symmetric form `I * (a * I^T * b)` is built -- all
+      factors are gathered at the highest-priority "block" location via
+      `I^T`, multiplied there, and the product is interpolated to `func`
+      via `I`. Use this for operators whose continuous form decomposes as
+      `I * A * I^T` (e.g. the elastic stiffness `σ = C ε`).
+
+    See `examples/userapi/08_staggered_interp.ipynb` for the maths and a
+    worked elastic-stiffness example.
     """
 
     DERIV_COLLECT = True
diff --git a/devito/finite_differences/derivative.py b/devito/finite_differences/derivative.py
@@ -89,7 +89,14 @@ class Derivative(sympy.Derivative, Differentiable, Pickable):
     evaluation are `x0`, `fd_order` and `side`.
     """
 
-    _fd_priority = .9
+    @cached_property
+    def _fd_priority(self):
+        # A Derivative inherits the priority of its underlying expression, so
+        # that `highest_priority(C*v.dx)` and `highest_priority((C*v).dx)`
+        # agree on the gather location and the two gathering paths
+        # (`_gather_for_diff` and `Mul._eval_at(interp_mode='symmetric')`)
+        # produce consistent answers.
+        return getattr(self.expr, '_fd_priority', 0)
 
     __rargs__ = ('expr', '*dims')
     __rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs',
@@ -472,7 +479,7 @@ def T(self):
 
         return self._rebuild(transpose=adjoint)
 
-    def _eval_at(self, func, mul_first=False, **kwargs):
+    def _eval_at(self, func, interp_mode='direct', **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
@@ -522,7 +529,7 @@ def _eval_at(self, func, mul_first=False, **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, mul_first=mul_first, **kwargs)
+            args = [self._rebuild(a)._eval_at(func, interp_mode=interp_mode, **kwargs)
                     for a in args]
             return self.expr.func(*args)
         elif self.expr.is_Mul:
@@ -594,7 +601,7 @@ def _eval_fd(self, expr, **kwargs):
             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=x0_deriv, expand=expand)
+                                     x0=self.x0, expand=expand)
 
         # Step 4: Apply substitutions
         for e in self._ppsubs:
diff --git a/devito/finite_differences/differentiable.py b/devito/finite_differences/differentiable.py
@@ -1,4 +1,5 @@
 from collections import ChainMap
+from contextlib import suppress
 from functools import cached_property, singledispatch
 from itertools import product
 
@@ -185,8 +186,10 @@ def coefficients(self):
         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)
-                           for a in self.args])
+        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 _subs(self, old, new, **hints):
         if old == self:
@@ -576,9 +579,6 @@ def _fd_priority(self):
             return super()._fd_priority
         return highest_priority(self)._fd_priority
 
-    def _eval_at(self, func, **kwargs):
-        return self
-
 
 class Add(DifferentiableOp, sympy.Add):
     __sympy_class__ = sympy.Add
@@ -668,67 +668,63 @@ def _gather_for_diff(self):
             other = self.func(*other)._eval_at(highest_priority(self))
             return self.func(other, *derivs)
 
-    def _eval_at(self, func, mul_first=False, **kwargs):
-        # Dont evaluate mul first
-        if not mul_first:
-            return super()._eval_at(func, mul_first=mul_first)
-
-        # Same staggering, no need to interpolate
-        if self.staggered == func.staggered:
-            return self
-
-        # Get highest priority function for evaluation
-        func0 = highest_priority(self, ref=func)
-
-        # Not a basic a*b*c... expression, expand
-        if any(isinstance(f, DifferentiableOp) for f in self.args):
-            return diffify(self._eval_expand_mul())._eval_at(func, mul_first=mul_first)
-
-        # Split Derivative and Differentiable args
-        derivs, other = split(self.args, lambda e: isinstance(e, sympy.Derivative))
-
-        # Evaluate all at highest priority function
-        if derivs:
-            derivs = self.func(*[d._eval_at(func, mul_first=mul_first) for d in derivs])
-        else:
-            derivs = 1
-
-        if not other:
-            return derivs
-        expr = self.func(*other)
-
-        # Non differentiable expr (e.g., number)
-        if not isinstance(expr, Differentiable):
-            return self.func(derivs, expr)
-
-        # Evaluate expression at func_args
-        print(f"\nEvaluating expr {expr} at func0 {func0} for func {func} from {self}")
-        expr = Differentiable._eval_at(expr, func0, mul_first=False)
-
-        # Interpolate derivatives at func0
-        x0 = {d: v for d, v in func0.indices_ref.getters.items()
-            if not d.is_Time and v is not func.indices_ref.getters.get(d, d)}
-        if x0 and not derivs == 1:
-            print(f"Interpolating derivs {derivs} x0={derivs.x0} at {x0}")
-            derivs = derivs.diff(*x0.keys(), deriv_order=(0,)*len(x0),
-                                fd_order=(self.interp_order,)*len(x0),
-                                x0=x0)
-        newexpr = self.func(derivs, expr)
-
-        # Finally at func
-        if not func.staggered == func0.staggered:
-            x0_f = {d: v for d, v in func.indices_ref.getters.items()
-                    if not d.is_Time and v is not func0.indices_ref.getters.get(d)}
-            if x0_f:
-                print(f"Final interpolation of derivs {self.func(derivs, expr)} at func {x0_f}")
-                return newexpr.diff(*x0_f.keys(), deriv_order=(0,)*len(x0_f),
-                                    fd_order=(self.interp_order,)*len(x0_f),
-                                    x0=x0_f)
+    def _eval_at(self, func, interp_mode='direct', **kwargs):
+        """
+        Evaluate a Mul at the location of `func`.
+
+        Two modes:
+
+        - `interp_mode='direct'` (default): per-arg evaluation; each factor is
+          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.
+            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`.
+        """
+        if interp_mode != 'symmetric':
+            return super()._eval_at(func, **kwargs)
+
+        diff_args = [a for a in self.args if isinstance(a, Differentiable)]
+        other_args = [a for a in self.args if not isinstance(a, Differentiable)]
+
+        # Symmetric form requires every Differentiable factor to differ from
+        # func; otherwise direct evaluation is cleaner and equivalent.
+        if len(diff_args) < 2 or \
+           any(a.staggered == func.staggered for a in diff_args):
+            return super()._eval_at(func, **kwargs)
+
+        block_indices = highest_priority(self).indices_ref
+
+        # Bring each factor to block's location (I^T where needed)
+        new_factors = list(other_args)
+        for a in diff_args:
+            if isinstance(a, sympy.Derivative):
+                source = _post_x0_indices(a, func)
+                a = a._rebuild(x0={dim: func.indices_ref[dim] for dim in a.dims
+                                   if dim in func.indices_ref.getters})
             else:
-                return newexpr
-        else:
-            # Return the full expression with Derivatives
-            return newexpr
+                source = a.indices_ref
+            new_factors.append(_interp_at(a, source, block_indices,
+                                          self.interp_order))
+
+        # Final I from block's location to func
+        return _interp_at(self.func(*new_factors), block_indices,
+                          func.indices_ref, self.interp_order)
 
 
 class Pow(DifferentiableOp, sympy.Pow):
@@ -1255,6 +1251,63 @@ def _diff2sympy(obj):
 evalf_table[Pow] = evalf_table[sympy.Pow]
 
 
+def _interp_mapper(source, target, dims):
+    """
+    Build a `{dim: target_index}` mapper for dimensions in `dims` where
+    `source[dim]` differs from `target[dim]`.
+
+    `source` and `target` are dict-like `{dim: index_expr}` (e.g. a plain
+    dict or a `DimensionTuple`). Dimensions missing from either side are
+    skipped silently.
+    """
+    mapper = {}
+    for d in dims:
+        try:
+            s = source[d]
+            t = target[d]
+        except (KeyError, IndexError):
+            continue
+        if s is not t:
+            mapper[d] = t
+    return mapper
+
+
+def _interp_at(expr, source, target, interp_order):
+    """
+    Build a symbolic 0-order FD interpolation operator on `expr` that maps
+    values from `source` indices to `target` indices, only on the
+    dimensions where the two locations differ.
+    """
+    if not isinstance(expr, Differentiable):
+        return expr
+
+    mapper = _interp_mapper(source, target, expr.dimensions)
+    if not mapper:
+        return expr
+
+    return expr.diff(*mapper.keys(),
+                     deriv_order=(0,) * len(mapper),
+                     fd_order=(interp_order,) * len(mapper),
+                     x0=mapper)
+
+
+def _post_x0_indices(deriv, func):
+    """
+    Conceptual indices of `deriv` after setting `x0` on its own derivative
+    dimensions to `func`'s indices. Derivative dims take `func`'s indices;
+    other dims keep the underlying expression's natural location (so that
+    `interp_for_fd` does not introduce a spurious second shift).
+    """
+    ref = {}
+    for dim in deriv.dimensions:
+        if dim in deriv.dims and dim in func.indices_ref.getters:
+            ref[dim] = func.indices_ref[dim]
+        else:
+            with suppress(KeyError):
+                ref[dim] = deriv.indices_ref[dim]
+    return ref
+
+
 # Interpolation for finite differences
 @singledispatch
 def interp_for_fd(expr, x0, **kwargs):
diff --git a/devito/operator/operator.py b/devito/operator/operator.py
@@ -341,7 +341,7 @@ def _lower_exprs(cls, expressions, **kwargs):
             * Shift indices for domain alignment.
         """
         expand = kwargs['options'].get('expand', True)
-        mul_first = kwargs['options'].get('eval-mul-first', False)
+        interp_mode = kwargs['options'].get('interp-mode', 'direct')
 
         # Specialization is performed on unevaluated expressions
         expressions = cls._specialize_dsl(expressions, **kwargs)
@@ -352,7 +352,7 @@ def _lower_exprs(cls, expressions, **kwargs):
         # ModuloDimensions
         if not expand:
             expand = lambda d: d.is_Stepping
-        expressions = flatten([i._evaluate(expand=expand, mul_first=mul_first)
+        expressions = flatten([i._evaluate(expand=expand, interp_mode=interp_mode)
                                for i in expressions])
 
         # Scalarize the tensor equations, if any
diff --git a/devito/types/dense.py b/devito/types/dense.py
@@ -15,6 +15,7 @@
 from devito.deprecations import deprecations
 from devito.exceptions import InvalidArgument
 from devito.finite_differences import Differentiable, generate_fd_shortcuts
+from devito.finite_differences.differentiable import _interp_mapper
 from devito.finite_differences.tools import fd_weights_registry
 from devito.logger import debug, warning
 from devito.mpi import MPI
@@ -1116,22 +1117,24 @@ def __fd_setup__(self):
 
     @cached_property
     def _fd_priority(self):
+        # NODE takes precedence: coefficients are conventionally stored at the
+        # cell centre, so when we gather a product onto a single location
+        # (either via _gather_for_diff or symmetric Mul._eval_at), NODE is the
+        # natural one to pick.
         return 1.2 if self.staggered.on_node else 1.1
 
     def _eval_at(self, func, **kwargs):
         if self.staggered == func.staggered or self.interp_order == 0:
             return self
 
-        mapper = {}
-        for d in self.dimensions:
-            try:
-                if self.indices_ref[d] is not func.indices_ref[d]:
-                    f_idx = func.indices_ref[d]._subs(func.dimensions[d], d)
-                    mapper[self.indices_ref[d]] = f_idx
-            except KeyError:
-                pass
+        # Dims where self and func indices differ -> {dim: func_idx}
+        diff = _interp_mapper(self.indices_ref, func.indices_ref, self.dimensions)
+
+        # Translate into a subs mapper {self_idx: func_idx} aligned on self's dims
+        subs_map = {self.indices_ref[d]: t._subs(func.dimensions[d], d)
+                    for d, t in diff.items()}
 
-        return self.subs(mapper)
+        return self.subs(subs_map)
 
     @classmethod
     def __staggered_setup__(cls, dimensions, staggered=None, **kwargs):
diff --git a/devito/types/utils.py b/devito/types/utils.py
@@ -58,12 +58,14 @@ def on_node(self):
         return not self or all(s == 0 for s in self)
 
     def __eq__(self, other):
-        if not isinstance(other, Staggering):
-            return False
-        all_same = self and other and all(a == b for a, b in zip(self, other))
-        all_node = self.on_node and other.on_node
+        # Two empty-or-all-zero Staggerings are equivalent regardless of arity
+        # (a Function declared with `staggered=NODE` and one declared without
+        # both live at the cell centre).
+        if isinstance(other, Staggering) and self.on_node and other.on_node:
+            return True
+        return tuple.__eq__(self, other)
 
-        return all_same or all_node
+    __hash__ = DimensionTuple.__hash__
 
     @property
     def _ref(self):
@@ -74,8 +76,6 @@ def _ref(self):
         else:
             return tuple(d for d, s in zip(self.getters, self, strict=True) if s == 1)
 
-    __hash__ = DimensionTuple.__hash__
-
 
 class IgnoreDimSort(tuple):
     """A tuple subclass used to wrap the implicit_dims to indicate
diff --git a/examples/userapi/08_staggered_interp.ipynb b/examples/userapi/08_staggered_interp.ipynb
diff --git a/tests/test_differentiable.py b/tests/test_differentiable.py
diff --git a/tests/test_staggered_utils.py b/tests/test_staggered_utils.py
