api: prevent evaluated derivatives to be re-evaluted

Author: mloubout

devito/finite_differences/differentiable.py

@@ -952,6 +952,11 @@ def _evaluate(self, **kwargs):
 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
 for i in ('DiffDerivative', 'IndexDerivative'):
@@ -998,6 +1003,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

@@ -995,9 +995,23 @@ 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
+                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:
                     continue
@@ -1008,6 +1022,11 @@ def _grid_map(self):
                     mapper.update({d: i})
             except (AttributeError, TypeError):
                 mapper.update({d: i})
+
+        # Substitutions for self.function
+        if mapper:
+            mapper['subs'] = subs
+
         return mapper
 
     def _evaluate(self, **kwargs):
@@ -1019,29 +1038,33 @@ 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)
+        if subs:
+            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

@@ -1073,9 +1073,15 @@ 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]}
+
+        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
+
         if mapper:
             return self.subs(mapper)
         return self

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

@@ -716,7 +716,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))