Merge pull request #2754 from devitocodes/interp-order

api: add interpolation order api for staggered/off grid evaluation

Author: mloubout

devito/finite_differences/differentiable.py

@@ -36,7 +36,7 @@ class Differentiable(sympy.Expr, Evaluable):
     # operators to be used
     _op_priority = sympy.Expr._op_priority + 1.
 
-    __rkwargs__ = ('space_order', 'time_order', 'indices')
+    __rkwargs__ = ('space_order', 'interp_order', 'time_order', 'indices')
 
     @cached_property
     def _functions(self):
@@ -54,6 +54,12 @@ def space_order(self):
         return min([getattr(i, 'space_order', 100) or 100 for i in self._args_diff],
                    default=100)
 
+    @cached_property
+    def interp_order(self):
+        # Default 2 is a reasonable default for interpolation
+        return min([getattr(i, 'interp_order', 2) or 2 for i in self._args_diff],
+                   default=2)
+
     @cached_property
     def time_order(self):
         # Default 100 is for "infinitely" differentiable
@@ -1129,7 +1135,7 @@ def _(expr, x0, **kwargs):
     x0_expr = {d: v for d, v in x0.items() if v is not expr.indices_ref[d]}
     if x0_expr:
         dims = tuple((d, 0) for d in x0_expr)
-        fd_o = tuple([2]*len(dims))
+        fd_o = tuple([expr.interp_order]*len(dims))
         return Derivative(expr, *dims, fd_order=fd_o, x0=x0_expr)
     else:
         return expr

devito/types/basic.py

@@ -1023,6 +1023,7 @@ def _evaluate(self, **kwargs):
         if not self._grid_map:
             return self
 
+        io = self.interp_order
         # Base function
         if self._avg_mode == 'harmonic':
             retval = 1 / self.function
@@ -1031,7 +1032,7 @@ def _evaluate(self, **kwargs):
 
         # Apply interpolation from inner most dim
         for d, i in self._grid_map.items():
-            retval = retval.diff(d, deriv_order=0, fd_order=2, x0={d: i})
+            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

devito/types/dense.py

@@ -941,6 +941,10 @@ class Function(DiscreteFunction):
             discretization order (`o`) as well as the number of points on
             the left/right sides of a generic point of interest for each
             SpaceDimension.
+    interp_order: int, optional, default=2
+        Order of the interpolation scheme used to evaluate the Function at
+        non-grid points (e.g., when using a Function as a parameter to be
+        evaluated at a staggered location).
     shape : tuple of ints, optional
         Shape of the domain region in grid points. Only necessary if `grid`
         isn't given.
@@ -1014,7 +1018,7 @@ class Function(DiscreteFunction):
     is_autopaddable = True
 
     __rkwargs__ = (DiscreteFunction.__rkwargs__ +
-                   ('space_order', 'dimensions'))
+                   ('space_order', 'interp_order', 'dimensions'))
 
     def _cache_meta(self):
         # Attach additional metadata to self's cache entry
@@ -1036,6 +1040,16 @@ def __init_finalize__(self, *args, **kwargs):
         else:
             raise TypeError("Invalid `space_order`")
 
+        # Interpolation order
+        interp_order = kwargs.get('interp_order', 2)
+        if not is_integer(interp_order):
+            raise TypeError("`interp_order` must be an integer")
+        elif interp_order < 1:
+            raise ValueError("`interp_order` must be >= 2")
+        elif interp_order > self._space_order and self._space_order > 1:
+            raise ValueError("`interp_order` must be <= `space_order`")
+        self._interp_order = interp_order
+
         # Acquire derivative shortcuts
         if self is self.function:
             self._fd = self.__fd_setup__()
@@ -1233,6 +1247,11 @@ def space_order(self):
         """The space order."""
         return self._space_order
 
+    @property
+    def interp_order(self):
+        """The interpolation order."""
+        return self._interp_order
+
     def sum(self, p=None, dims=None):
         """
         Generate a symbolic expression computing the sum of `p` points

examples/userapi/01_dsl.ipynb

@@ -2,6 +2,7 @@
 
 import sympy
 import pytest
+import numpy as np
 
 from devito import Function, Grid, Differentiable, NODE
 from devito.finite_differences.differentiable import (Add, Mul, Pow, diffify,
@@ -96,28 +97,51 @@ def sp_diff(a, b):
 
 
 @pytest.mark.parametrize('ndim', [1, 2, 3])
-def test_avg_mode(ndim):
+@pytest.mark.parametrize('io', [None, 2, 4])
+def test_avg_mode(ndim, io):
     grid = Grid([11]*ndim)
-    v = Function(name='v', grid=grid, staggered=grid.dimensions)
-    a0 = Function(name="a0", grid=grid)
-    a = Function(name="a", grid=grid, parameter=True)
-    b = Function(name="b", grid=grid, parameter=True, avg_mode='harmonic')
+    v = Function(name='v', grid=grid, staggered=grid.dimensions, space_order=4)
+    kw = {'space_order': 4}
+    if io is not None:
+        kw['interp_order'] = io
+    else:
+        io = 2  # Default value
+
+    with pytest.raises(ValueError):
+        # interp_order > space_order
+        Function(name="a", grid=grid, parameter=True, interp_order=8, space_order=4)
+    with pytest.raises(ValueError):
+        # interp_order < 1
+        Function(name="a", grid=grid, parameter=True, interp_order=0, space_order=4)
+    with pytest.raises(TypeError):
+        # interp_order not int
+        Function(name="a", grid=grid, parameter=True, interp_order=2.5, space_order=4)
+
+    a0 = Function(name="a0", grid=grid, **kw)
+    a = Function(name="a", grid=grid, parameter=True, **kw)
+    b = Function(name="b", grid=grid, parameter=True, avg_mode='harmonic', **kw)
 
     a0_avg = a0._eval_at(v)
-    a_avg = a._eval_at(v).evaluate
-    b_avg = b._eval_at(v).evaluate
+    a_avg = a._eval_at(v).evaluate.simplify()
+    b_avg = b._eval_at(v).evaluate.simplify()
 
     assert a0_avg == a0
 
     # Indices around the point at the center of a cell
-    all_shift = tuple(product(*[[0, 1] for _ in range(ndim)]))
+    idx = list(range(-io//2 + 1, io//2 + 1))
+    all_shift = tuple(product(*[idx for _ in range(ndim)]))
+    coeffs = {2: [0.5, 0.5], 4: [-1/16, 9/16, 9/16, -1/16]}[io]
+    vars = ['i', 'j', 'k'][:ndim]
+    rule = ','.join(vars) + '->' + ''.join(vars)
+    ndcoeffs = np.einsum(rule, *([coeffs]*ndim))
     args = [{d: d + i * d.spacing for d, i in zip(grid.dimensions, s)} for s in all_shift]
 
     # Default is arithmetic average
-    assert sympy.simplify(a_avg - 0.5**ndim * sum(a.subs(arg) for arg in args)) == 0
+    expected = sum(c * a.subs(arg) for c, arg in zip(ndcoeffs.flatten(), args))
+    assert sympy.simplify(a_avg - expected) == 0
 
     # Harmonic average, h(a[.5]) = 1/(.5/a[0] + .5/a[1])
-    expected = 1/(0.5**ndim * sum(1/b.subs(arg) for arg in args))
-    assert sympy.simplify(1/b_avg.args[0] - expected) == 0
+    expected = (sum(c / b.subs(arg) for c, arg in zip(ndcoeffs.flatten(), args)))
+    assert sympy.simplify(b_avg.args[0] - expected) == 0
     assert isinstance(b_avg, SafeInv)
     assert b_avg.base == b