Merge pull request #2788 from devitocodes/strict-mul-eval-at

api: add mul interp mode

Author: mloubout

devito/core/cpu.py

@@ -111,6 +111,7 @@ def _normalize_kwargs(cls, **kwargs):
             )
 
         kwargs['options'].update(o)
+        kwargs['sym_options'] = cls._normalize_sym_kwargs(**kwargs)
 
         return kwargs
 

devito/core/gpu.py

@@ -121,6 +121,7 @@ def _normalize_kwargs(cls, **kwargs):
             )
 
         kwargs['options'].update(o)
+        kwargs['sym_options'] = cls._normalize_sym_kwargs(**kwargs)
 
         return kwargs
 

devito/core/operator.py

@@ -171,6 +171,30 @@ class BasicOperator(Operator):
     The target language constructor, to be specified by subclasses.
     """
 
+    # ------------------------------------------------------------------
+    # Symbolic-level option defaults (`sym_opt`).
+    # These steer mathematical choices made during expression lowering,
+    # *not* code generation or performance. They are kept separate from
+    # the `opt` options above to keep the two concerns distinct.
+    # ------------------------------------------------------------------
+
+    INTERP_MODE = 'direct'
+    """
+    Default for the `sym_opt={'interp-mode': ...}` option. Controls how
+    a product of fields living at different staggered locations is mapped
+    onto a target location:
+
+    * `'direct'` (default): each factor is interpolated to the target
+      independently. Cheapest stencil.
+    * `'symmetric'`: factors are first gathered at a common "block"
+      location, multiplied there, and the result is interpolated once to
+      the target. Preserves the `I A I^T` matrix structure, so the
+      discrete operator stays self-adjoint when the continuous one is
+      (e.g. the elastic stiffness `sigma = C eps`).
+
+    See `examples/userapi/08_staggered_interp.ipynb` for a worked example.
+    """
+
     @classmethod
     def _normalize_kwargs(cls, **kwargs):
         # Will be populated with dummy values; this method is actually overridden
@@ -188,12 +212,30 @@ def _normalize_kwargs(cls, **kwargs):
             )
 
         kwargs['options'].update(o)
+        kwargs['sym_options'] = cls._normalize_sym_kwargs(**kwargs)
 
         return kwargs
 
+    @classmethod
+    def _normalize_sym_kwargs(cls, **kwargs):
+        """
+        Fill in defaults and validate keys for the `sym_opt` dict passed to
+        the Operator. Returns the normalized `sym_options` dict.
+        """
+        so = dict(kwargs.get('sym_options', {}))
+        out = {'interp-mode': so.pop('interp-mode', cls.INTERP_MODE)}
+
+        if so:
+            raise InvalidOperator(
+                f'Unrecognized symbolic options: [{", ".join(list(so))}]'
+            )
+
+        return out
+
     @classmethod
     def _check_kwargs(cls, **kwargs):
         oo = kwargs['options']
+        so = kwargs['sym_options']
 
         if oo['mpi'] and oo['mpi'] not in cls.MPI_MODES:
             raise InvalidOperator(f"Unsupported MPI mode `{oo['mpi']}`")
@@ -209,6 +251,9 @@ def _check_kwargs(cls, **kwargs):
         if oo['errctl'] not in (None, False, 'basic', 'max'):
             raise InvalidOperator("Illegal `errctl` value")
 
+        if so['interp-mode'] not in ('direct', 'symmetric'):
+            raise InvalidOperator("Illegal `interp-mode` value")
+
     def _autotune(self, args, setup):
         if setup in [False, 'off']:
             return args

devito/finite_differences/derivative.py

@@ -89,7 +89,14 @@ class Derivative(sympy.Derivative, Differentiable, Pickable):
     evaluation are `x0`, `fd_order` and `side`.
     """
 
-    _fd_priority = 3
+    @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):
+    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
@@ -525,7 +532,8 @@ def _eval_at(self, func):
                 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) for a in args]
+            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:
             # For Mul, We treat the basic case `u(x + h_x/2) * v(x) which is what appear

devito/finite_differences/differentiable.py

@@ -15,6 +15,7 @@
     # 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.tools import coeff_priority, make_shift_x0
 from devito.logger import warning
 from devito.tools import (
@@ -140,10 +141,6 @@ def indices_ref(self):
             return DimensionTuple(*self.dimensions, getters=self.dimensions)
         return highest_priority(self).indices_ref
 
-    @cached_property
-    def is_Staggered(self):
-        return any([getattr(i, 'is_Staggered', False) for i in self._args_diff])
-
     @cached_property
     def is_TimeDependent(self):
         return any(i.is_Time for i in self.dimensions)
@@ -184,13 +181,11 @@ def coefficients(self):
         key = lambda x: coeff_priority.get(x, -1)
         return sorted(coefficients, key=key, reverse=True)[0]
 
-    def _eval_at(self, func):
-        if not func.is_Staggered:
-            # Cartesian grid, do no waste time
-            return self
+    def _eval_at(self, func, **kwargs):
         return self.func(*[
-            getattr(a, '_eval_at', lambda x: a)(func) for a in self.args  # noqa: B023
-        ])  # false positive
+            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:
@@ -669,6 +664,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, 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, 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)
+
+        block_indices = highest_priority(self).indices_ref
+
+        # Bring each factor to block's location (I^T where needed)
+        new_factors = list(other)
+        for a in diff:
+            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:
+                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):
     _fd_priority = 0
@@ -1020,7 +1072,7 @@ def _subs(self, old, new, **hints):
 
 class DiffDerivative(IndexDerivative, DifferentiableOp):
 
-    def _eval_at(self, func):
+    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
         return self
@@ -1074,7 +1126,7 @@ def _new_rawargs(self, *args, **kwargs):
         kwargs.pop('is_commutative', None)
         return self.func(*args, **kwargs)
 
-    def _eval_at(self, func):
+    def _eval_at(self, func, **kwargs):
         # An EvalDerivative must have already been evaluated at a valid x0
         # and should not be re-evaluated at a different location
         return self
@@ -1092,7 +1144,7 @@ class diffify:
 
     Notes
     -----
-    The name "diffify" stems from SymPy's "simpify", which has an analogous task --
+    The name "diffify" stems from SymPy's "simplify", which has an analogous task --
     converting all arguments into SymPy core objects.
     """
 
@@ -1240,7 +1292,7 @@ def _(expr, x0, **kwargs):
 @interp_for_fd.register(AbstractFunction)
 def _(expr, x0, **kwargs):
     x0_expr = {d: v for d, v in x0.items() if v.has(d)
-               and expr.indices[d] is not v}
+               and expr.indices.get(d, v) is not v}
     if x0_expr:
         return expr.subs({expr.indices[d]: v for d, v in x0_expr.items()})
     else:

devito/finite_differences/interpolation.py

@@ -0,0 +1,62 @@
+from contextlib import suppress
+
+__all__ = ['interp_at', 'interp_mapper', 'post_x0_indices']
+
+
+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.
+    """
+    from devito.finite_differences.differentiable import Differentiable
+
+    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.dimensions:
+            ref[dim] = func.indices_ref[dim]
+        else:
+            with suppress(KeyError):
+                ref[dim] = deriv.indices_ref[dim]
+    return ref

devito/finite_differences/tools.py

@@ -292,7 +292,7 @@ def generate_indices(expr, dim, order, side=None, matvec=None, x0=None, nweights
     o_min = int(np.ceil(mid - r)) + side.val
     o_max = int(np.floor(mid + r)) + side.val
     if o_max == o_min:
-        if dim.is_Time or not expr.is_Staggered:
+        if dim.is_Time or not bool(expr.staggered):
             o_max += 1
         else:
             o_min -= 1

devito/operator/operator.py

@@ -68,6 +68,16 @@ class Operator(Callable):
             Symbolic substitutions to be applied to ``expressions``.
         * opt : str
             The performance optimization level. Defaults to ``configuration['opt']``.
+        * sym_opt : dict
+            Symbolic-level options controlling mathematical choices made during
+            expression lowering (e.g. how staggered multi-factor products are
+            interpolated). Distinct from ``opt``, which controls code generation
+            and performance. Accepted keys:
+
+            - ``'interp-mode'`` (``'direct'`` | ``'symmetric'``): selects the
+              interpolation strategy used by ``Mul._eval_at`` when projecting a
+              multi-factor expression onto a target staggered location. See the
+              tutorial at ``examples/userapi/08_staggered_interp.ipynb``.
         * language : str
             The target language for shared-memory parallelism. Defaults to
             ``configuration['language']``.
@@ -235,6 +245,7 @@ def _build(cls, expressions, **kwargs):
         # Potentially required for lazily allocated Functions
         op._mode = kwargs['mode']
         op._options = kwargs['options']
+        op._sym_options = kwargs['sym_options']
         op._allocator = kwargs['allocator']
         op._platform = kwargs['platform']
 
@@ -342,6 +353,7 @@ def _lower_exprs(cls, expressions, **kwargs):
             * Shift indices for domain alignment.
         """
         expand = kwargs['options'].get('expand', True)
+        interp_mode = kwargs.get('sym_options', {}).get('interp-mode', 'direct')
 
         # Specialization is performed on unevaluated expressions
         expressions = cls._specialize_dsl(expressions, **kwargs)
@@ -352,7 +364,8 @@ def _lower_exprs(cls, expressions, **kwargs):
         # ModuloDimensions
         if not expand:
             expand = lambda d: d.is_Stepping
-        expressions = flatten([i._evaluate(expand=expand) for i in expressions])
+        expressions = flatten([i._evaluate(expand=expand, interp_mode=interp_mode)
+                               for i in expressions])
 
         # Scalarize the tensor equations, if any
         expressions = [j for i in expressions for j in i._flatten]
@@ -1661,6 +1674,12 @@ def parse_kwargs(**kwargs):
         mode = 'noop'
     kwargs['mode'] = mode
 
+    # `sym_opt` -- symbolic-level options (mathematical choices, not codegen)
+    sym_opt = kwargs.pop('sym_opt', None) or {}
+    if not isinstance(sym_opt, (dict, frozendict)):
+        raise InvalidOperator(f"Illegal `sym_opt={str(sym_opt)}`")
+    kwargs['sym_options'] = dict(sym_opt)
+
     # `platform`
     platform = kwargs.get('platform')
     if platform is not None: