tweak mul mode

Author: mloubout

devito/finite_differences/derivative.py

@@ -89,7 +89,7 @@ class Derivative(sympy.Derivative, Differentiable, Pickable):
     evaluation are `x0`, `fd_order` and `side`.
     """
 
-    _fd_priority = 3
+    _fd_priority = .9
 
     __rargs__ = ('expr', '*dims')
     __rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs',
@@ -472,7 +472,7 @@ def T(self):
 
         return self._rebuild(transpose=adjoint)
 
-    def _eval_at(self, func, **kwargs):
+    def _eval_at(self, func, mul_first=False, **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 +522,8 @@ def _eval_at(self, func, **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, **kwargs) for a in args]
+            args = [self._rebuild(a)._eval_at(func, mul_first=mul_first, **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

@@ -185,17 +185,11 @@ def coefficients(self):
         return sorted(coefficients, key=key, reverse=True)[0]
 
     def _eval_at(self, func, **kwargs):
-        if not func.is_Staggered:
+        if not func.is_Staggered and not self.is_Staggered:
             # Cartesian grid, do no waste time
             return self
-<<<<<<< HEAD
-        return self.func(*[
-            getattr(a, '_eval_at', lambda x: a)(func) for a in self.args  # noqa: B023
-        ])  # false positive
-=======
         return self.func(*[getattr(a, '_eval_at', lambda x, **kw: a)(func, **kwargs)
                            for a in self.args])
->>>>>>> 198aedce4 (api: add mul interp mode)
 
     def _subs(self, old, new, **hints):
         if old == self:
@@ -497,14 +491,19 @@ def has_free(self, *patterns):
             return all(i in self.free_symbols for i in patterns)
 
 
+<<<<<<< HEAD
 def highest_priority(diff_op):
     if not diff_op._args_diff:
         return diff_op
 
+=======
+def highest_priority(DiffOp, ref=None):
+>>>>>>> c34bae6f0 (tweak mul mode)
     # 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
+<<<<<<< HEAD
     prio = lambda x: (getattr(x, '_fd_priority', 0), len(x.dimensions))
 <<<<<<< HEAD
     prio_func = sorted(diff_op._args_diff, key=prio, reverse=True)[0]
@@ -514,6 +513,14 @@ def highest_priority(diff_op):
         return highest_priority(prio_func)
     return prio_func
 =======
+=======
+    def stagg(x):
+        try:
+            return int(x.staggered == ref.staggered)
+        except AttributeError:
+            return 0
+    prio = lambda x: (stagg(x), getattr(x, '_fd_priority', 0), len(x.dimensions))
+>>>>>>> c34bae6f0 (tweak mul mode)
     args = DiffOp._args_diff
     if not args:
         return DiffOp
@@ -692,50 +699,62 @@ def _eval_at(self, func, mul_first=False, **kwargs):
         if not mul_first:
             return super()._eval_at(func, mul_first=mul_first)
 
-        # Not a basic a*b*c... expression, just defer to superclass
+        # 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 super()._eval_at(func, mul_first=mul_first)
+            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 = Differentiable._eval_at(self.func(*derivs), func,
-                                             mul_first=mul_first)
+            derivs = self.func(*[d._eval_at(func, mul_first=mul_first) for d in derivs])
         else:
             derivs = 1
 
         if not other:
             return derivs
-        elif len(other) > 1:
-            expr = self.func(*other)._gather_for_diff
-        else:
-            expr = other[0]
+        expr = self.func(*other)
 
         # Non differentiable expr (e.g., number)
         if not isinstance(expr, Differentiable):
             return self.func(derivs, expr)
 
-        # Build mapper for dimensions that need to be interpolated
-        mapper = {}
-        for d in self.dimensions:
-            try:
-                if self.indices_ref[d] is not func.indices_ref[d]:
-                    mapper[d] = func.indices_ref[d]
-            except KeyError:
-                pass
-
-        # Nothing to interpolate
-        if not mapper:
-            return super()._eval_at(func, mul_first=mul_first)
-
-        # Interpolate expr at the required indices
-        interp = expr.diff(*mapper.keys(), deriv_order=[0 for _ in mapper],
-                           fd_order=[self.interp_order for _ in mapper],
-                           x0=mapper)
-
-        # Return the full expression with Derivatives
-        return self.func(derivs, interp)
+        # 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)
+            else:
+                return newexpr
+        else:
+            # Return the full expression with Derivatives
+            return newexpr
 
 
 class Pow(DifferentiableOp, sympy.Pow):

devito/types/utils.py

@@ -57,14 +57,15 @@ class Staggering(DimensionTuple):
     def on_node(self):
         return not self or all(s == 0 for s in self)
 
-    @property
-    def _ref(self):
-        if not self:
-            return None
-        elif self.on_node:
-            return NODE
-        else:
-            return tuple(d for d, s in zip(self.getters, self, strict=True) if s == 1)
+    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
+
+        return all_same or all_node
+
+    __hash__ = DimensionTuple.__hash__
 
 
 class IgnoreDimSort(tuple):