api: prevent evaluated derivatives to be re-evaluted

Author: mloubout

devito/finite_differences/derivative.py

@@ -335,20 +335,6 @@ def __call__(self, x0=None, fd_order=None, side=None, method=None, **kwargs):
             except AttributeError:
                 raise TypeError("fd_order incompatible with dimensions") from None
 
-        if isinstance(self.expr, Derivative):
-            # In case this was called on a perfect cross-derivative `u.dxdy`
-            # we need to propagate the call to the nested derivative
-            rkwe = dict(rkw)
-            rkwe.pop('weights', None)
-            if 'x0' in rkwe:
-                rkwe['x0'] = self._filter_dims(self.expr._filter_dims(rkw['x0']),
-                                               neg=True)
-            if fd_order is not None:
-                fdo = self.expr._filter_dims(_fd_order)
-                if fdo:
-                    rkwe['fd_order'] = fdo
-            rkw['expr'] = self.expr(**rkwe)
-
         if fd_order is not None:
             rkw['fd_order'] = self._filter_dims(_fd_order, as_tuple=True)
 

devito/finite_differences/differentiable.py

@@ -964,7 +964,11 @@ def _subs(self, old, new, **hints):
 
 
 class DiffDerivative(IndexDerivative, DifferentiableOp):
-    pass
+
+    def _eval_at(self, func):
+        # 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
 
 
 # SymPy args ordering is the same for Derivatives and IndexDerivatives
@@ -1012,6 +1016,11 @@ def _new_rawargs(self, *args, **kwargs):
         kwargs.pop('is_commutative', None)
         return self.func(*args, **kwargs)
 
+    def _eval_at(self, func):
+        # An EvalDerivative must have already been evaluated at a valid x0
+        # and should not be re-evaluated at a different location
+        return self
+
 
 class diffify:
 

devito/types/basic.py

@@ -984,25 +984,48 @@ def c0(self):
     def _eval_deriv(self):
         return self
 
-    @property
+    @cached_property
     def _grid_map(self):
         """
         Mapper of off-grid interpolation points indices for each dimension.
         """
         mapper = {}
+        subs = {}
         for i, j, d in zip(self.indices, self.indices_ref, self.dimensions):
             # Two indices are aligned if they differ by an Integer*spacing.
-            v = (i - j)/d.spacing
+            if not i.has(d):
+                # Maybe a SubDimension
+                dims = {sd for sd in i.free_symbols if getattr(sd, 'is_Dimension', False)
+                        and d in sd._defines}
+
+                # More than one Dimension, cannot handle
+                if len(dims) != 1:
+                    continue
+
+                # SubDimensions -> Dimension substitutions for interpolation
+                sd = dims.pop()
+                v = (i - j._subs(d, sd))/d.spacing
+                i = i._subs(sd, d)
+                subs[d] = sd
+            else:
+                v = (i - j)/d.spacing
+
             try:
                 if not isinstance(v, sympy.Number) or int(v) == v:
+                    # Skip if index is on grid
                     continue
-                # Skip if index is just a Symbol or integer
                 elif (i.is_Symbol and not i.has(d)) or i.is_Integer:
+                    # Skip if index is just a Symbol or integer
                     continue
                 else:
                     mapper.update({d: i})
             except (AttributeError, TypeError):
                 mapper.update({d: i})
+
+        # Substitutions for SubDimensions
+        if mapper:
+            mapper['subs'] = subs
+
         return mapper
 
     def _evaluate(self, **kwargs):
@@ -1014,29 +1037,32 @@ def _evaluate(self, **kwargs):
         This allow to evaluate off grid points as EvalDerivative that are better
         for the compiler.
         """
+        mapper = self._grid_map
+        subs = mapper.pop('subs', {})
         # Average values if at a location not on the Function's grid
-        if not self._grid_map:
+        if not mapper:
             return self
 
         io = self.interp_order
-        # Base function
         if self._avg_mode == 'harmonic':
-            retval = 1 / self.function
+            retval = 1 / self
         else:
-            retval = self.function
+            retval = self
 
         # Apply interpolation from inner most dim
-        for d, i in self._grid_map.items():
+        for d, i in mapper.items():
+            retval = retval._subs(i.subs(subs), self.indices_ref[d])
             retval = retval.diff(d, deriv_order=0, fd_order=io, x0={d: i})
 
         # Evaluate. Since we used `self.function` it will be on the grid when
         # evaluate is called again within FD
         retval = retval._evaluate(**kwargs)
+        retval = retval.subs(subs)
 
         # If harmonic averaging, invert at the end
         if self._avg_mode == 'harmonic':
             from devito.finite_differences.differentiable import SafeInv
-            retval = SafeInv(retval, self.function)
+            retval = SafeInv(retval, self.function.subs(subs))
 
         return retval
 

devito/types/dense.py

@@ -1075,12 +1075,16 @@ def _fd_priority(self):
     def _eval_at(self, func):
         if self.staggered == func.staggered:
             return self
-        mapper = {self.indices_ref[d]: func.indices_ref[d]
-                  for d in self.dimensions
-                  if self.indices_ref[d] is not func.indices_ref[d]}
-        if mapper:
-            return self.subs(mapper)
-        return self
+
+        mapper = {}
+        for d in self.dimensions:
+            try:
+                if self.indices_ref[d] is not func.indices_ref[d]:
+                    mapper[self.indices_ref[d]] = func.indices_ref[d]
+            except KeyError:
+                pass
+
+        return self.subs(mapper)
 
     @classmethod
     def __staggered_setup__(cls, dimensions, staggered=None, **kwargs):

examples/userapi/07_functions_on_subdomains.ipynb

@@ -3003,7 +3003,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "assert np.isclose(np.linalg.norm(rec.data), 4263.511, atol=0, rtol=1e-4)"
+    "assert np.isclose(np.linalg.norm(rec.data), 3640.584, atol=0, rtol=1e-4)"
    ]
   }
  ],

tests/test_derivatives.py

@@ -805,13 +805,11 @@ def test_param_stagg_add(self):
         eq1 = Eq(vx, (c11 * txx).dy)
         eq2 = Eq(vx, (c11 * txx + c66 * txy).dy)
 
-        # C66 is a paramater. Expects to evaluate c66 at xp then the derivative at yp
-        # and the derivative will interpolate txy at xp
+        # Expects to evaluate c66 at xp then the derivative at yp
         expect0 = (c66.subs({x: xp, y: yp}).evaluate * txy).dy.evaluate
         assert simplify(eq0.evaluate.rhs - expect0) == 0
 
-        # C11 is a paramater and txy is staggered in x.
-        # Expects to evaluate c11 and txy xp then the derivative at yp
+        # Expects to evaluate c11 and txy at xp then the derivative at yp
         expect1 = (c11._subs(x, xp).evaluate * txx._subs(x, xp).evaluate).dy.evaluate
         assert simplify(eq1.evaluate.rhs - expect1) == 0
 

tests/test_differentiable.py

@@ -57,7 +57,7 @@ def test_shift():
     assert a.shift(x, x.spacing).shift(x, -x.spacing) == a
     assert a.shift(x, x.spacing).shift(x, x.spacing) == a.shift(x, 2*x.spacing)
     assert a.dx.evaluate.shift(x, x.spacing) == a.shift(x, x.spacing).dx.evaluate
-    assert a.shift(x, .5 * x.spacing)._grid_map == {x: x + .5 * x.spacing}
+    assert a.shift(x, .5 * x.spacing)._grid_map == {x: x + .5 * x.spacing, 'subs': {}}
 
 
 def test_interp():

tests/test_mpi.py

@@ -2895,9 +2895,9 @@ def test_staggering(self, mode):
 
         op(time_M=2)
         # Expected norms computed "manually" from sequential runs
-        assert np.isclose(norm(ux), 6054.139, rtol=1.e-4)
-        assert np.isclose(norm(uxx), 17814.95, rtol=1.e-4)
-        assert np.isclose(norm(uxy), 58712.22, rtol=1.e-4)
+        assert np.isclose(norm(ux), 7003.098, rtol=1.e-4)
+        assert np.isclose(norm(uxx), 78902.21, rtol=1.e-4)
+        assert np.isclose(norm(uxy), 71852.62, rtol=1.e-4)
 
     @pytest.mark.parallel(mode=2)
     def test_op_new_dist(self, mode):

tests/test_symbolics.py

@@ -762,7 +762,7 @@ def test_is_on_grid():
     u = Function(name="u", grid=grid, space_order=2)
 
     assert u._grid_map == {}
-    assert u.subs({x: x0})._grid_map == {x: x0}
+    assert u.subs({x: x0})._grid_map == {x: x0, 'subs': {}}
     assert all(uu._grid_map == {} for uu in retrieve_functions(u.subs({x: x0}).evaluate))