Repository: devitocodes/devito
Issue: `div`, `grad` etc. should accept `shift` as a `tuple`

The tensor operators `div`, `grad` and `curl` currently only accept `shift` as a single `float`. Since the staggering can be different in each dimension however they should also be able to accept a `tuple` of `float` of length `dims`.

See, e.g., [here](https://github.com/devitocodes/devito/blob/1802de464921b2ca225ec5e18e1f7935a9fbddab/devito/types/tensor.py#L305).

Changed files:
tests/test_derivatives.py, tests/test_tensors.py, devito/finite_differences/differentiable.py, devito/finite_differences/tools.py, devito/types/tensor.py

diff --git a/devito/finite_differences/differentiable.py b/devito/finite_differences/differentiable.py
index 45784794e..51a4c12cd 100644
--- a/devito/finite_differences/differentiable.py
+++ b/devito/finite_differences/differentiable.py
@@ -7,6 +7,7 @@ from sympy.core.decorators import call_highest_priority
 from sympy.core.evalf import evalf_table
 
 from cached_property import cached_property
+from devito.finite_differences.tools import make_shift_x0
 from devito.logger import warning
 from devito.tools import filter_ordered, flatten
 from devito.types.lazy import Evaluable, EvalDerivative
@@ -242,15 +243,16 @@ class Differentiable(sympy.Expr, Evaluable):
 
     def div(self, shift=None):
         space_dims = [d for d in self.dimensions if d.is_Space]
-        return Add(*[getattr(self, 'd%s' % d.name)(x0=None if shift is None else
-                                                   d + shift * d.spacing)
-                     for d in space_dims])
+        shift_x0 = make_shift_x0(shift, (len(space_dims),))
+        return Add(*[getattr(self, 'd%s' % d.name)(x0=shift_x0(shift, d, None, i))
+                     for i, d in enumerate(space_dims)])
 
     def grad(self, shift=None):
         from devito.types.tensor import VectorFunction, VectorTimeFunction
-        comps = [getattr(self, 'd%s' % d.name)(x0=None if shift is None else
-                                               d + shift * d.spacing)
-                 for d in self.dimensions if d.is_Space]
+        space_dims = [d for d in self.dimensions if d.is_Space]
+        shift_x0 = make_shift_x0(shift, (len(space_dims),))
+        comps = [getattr(self, 'd%s' % d.name)(x0=shift_x0(shift, d, None, i))
+                 for i, d in enumerate(space_dims)]
         vec_func = VectorTimeFunction if self.is_TimeDependent else VectorFunction
         return vec_func(name='grad_%s' % self.name, time_order=self.time_order,
                         space_order=self.space_order, components=comps, grid=self.grid)
diff --git a/devito/finite_differences/tools.py b/devito/finite_differences/tools.py
index 3f517cdee..2b28d9c1c 100644
--- a/devito/finite_differences/tools.py
+++ b/devito/finite_differences/tools.py
@@ -5,7 +5,6 @@ import numpy as np
 from sympy import S, finite_diff_weights, cacheit, sympify
 
 from devito.tools import Tag, as_tuple
-from devito.finite_differences import Differentiable
 
 
 class Transpose(Tag):
@@ -47,13 +46,11 @@ centered = Side('centered', 0)
 def check_input(func):
     @wraps(func)
     def wrapper(expr, *args, **kwargs):
-        if expr.is_Number:
-            return S.Zero
-        elif not isinstance(expr, Differentiable):
-            raise ValueError("`%s` must be of type Differentiable (found `%s`)"
+        try:
+            return S.Zero if expr.is_Number else func(expr, *args, **kwargs)
+        except AttributeError:
+            raise ValueError("'%s' must be of type Differentiable, not %s"
                              % (expr, type(expr)))
-        else:
-            return func(expr, *args, **kwargs)
     return wrapper
 
 
@@ -261,3 +258,24 @@ def generate_indices_staggered(func, dim, order, side=None, x0=None):
             ind = [start, start - diff]
 
     return start, tuple(ind)
+
+
+def make_shift_x0(shift, ndim):
+    """
+    Returns a callable that calculates a shifted origin for each derivative
+    of an operation derivatives scheme (given by ndim) given a shift object
+    which can be a None, a float or a tuple with shape equal to ndim
+    """
+    if shift is None:
+        return lambda s, d, i, j: None
+    elif isinstance(shift, float):
+        return lambda s, d, i, j: d + s * d.spacing
+    elif type(shift) is tuple and np.shape(shift) == ndim:
+        if len(ndim) == 1:
+            return lambda s, d, i, j: d + s[j] * d.spacing
+        elif len(ndim) == 2:
+            return lambda s, d, i, j: d + s[i][j] * d.spacing
+        else:
+            raise ValueError("ndim length must be equal to 1 or 2")
+    raise ValueError("shift parameter must be one of the following options: "
+                     "None, float or tuple with shape equal to %s" % (ndim,))
diff --git a/devito/types/tensor.py b/devito/types/tensor.py
index b72e574a7..e35ebdc2b 100644
--- a/devito/types/tensor.py
+++ b/devito/types/tensor.py
@@ -7,6 +7,7 @@ from sympy.core.decorators import call_highest_priority
 from sympy.core.sympify import converter as sympify_converter
 
 from devito.finite_differences import Differentiable
+from devito.finite_differences.tools import make_shift_x0
 from devito.types.basic import AbstractTensor
 from devito.types.dense import Function, TimeFunction
 from devito.types.utils import NODE
@@ -218,10 +219,13 @@ class TensorFunction(AbstractTensor):
         Divergence of the TensorFunction (is a VectorFunction).
         """
         comps = []
-        func = vec_func(self, self)
-        for j, d in enumerate(self.space_dimensions):
+        func = vec_func(self)
+        ndim = len(self.space_dimensions)
+        shift_x0 = make_shift_x0(shift, (ndim, ndim))
+        for i in range(len(self.space_dimensions)):
             comps.append(sum([getattr(self[j, i], 'd%s' % d.name)
-                              for i, d in enumerate(self.space_dimensions)]))
+                              (x0=shift_x0(shift, d, i, j))
+                              for j, d in enumerate(self.space_dimensions)]))
         return func._new(comps)
 
     @property
@@ -230,7 +234,7 @@ class TensorFunction(AbstractTensor):
         Laplacian of the TensorFunction.
         """
         comps = []
-        func = vec_func(self, self)
+        func = vec_func(self)
         for j, d in enumerate(self.space_dimensions):
             comps.append(sum([getattr(self[j, i], 'd%s2' % d.name)
                               for i, d in enumerate(self.space_dimensions)]))
@@ -240,7 +244,7 @@ class TensorFunction(AbstractTensor):
         raise AttributeError("Gradient of a second order tensor not supported")
 
     def new_from_mat(self, mat):
-        func = tens_func(self, self)
+        func = tens_func(self)
         return func._new(self.rows, self.cols, mat)
 
     def classof_prod(self, other, mat):
@@ -306,8 +310,8 @@ class VectorFunction(TensorFunction):
         """
         Divergence of the VectorFunction, creates the divergence Function.
         """
-        return sum([getattr(self[i], 'd%s' % d.name)(x0=None if shift is None else
-                                                     d + shift * d.spacing)
+        shift_x0 = make_shift_x0(shift, (len(self.space_dimensions),))
+        return sum([getattr(self[i], 'd%s' % d.name)(x0=shift_x0(shift, d, None, i))
                     for i, d in enumerate(self.space_dimensions)])
 
     @property
@@ -315,7 +319,7 @@ class VectorFunction(TensorFunction):
         """
         Laplacian of the VectorFunction, creates the Laplacian VectorFunction.
         """
-        func = vec_func(self, self)
+        func = vec_func(self)
         comps = [sum([getattr(s, 'd%s2' % d.name) for d in self.space_dimensions])
                  for s in self]
         return func._new(comps)
@@ -334,16 +338,19 @@ class VectorFunction(TensorFunction):
         comp2 = getattr(self[0], derivs[2]) - getattr(self[2], derivs[0])
         comp3 = getattr(self[1], derivs[0]) - getattr(self[0], derivs[1])
 
-        vec_func = VectorTimeFunction if self.is_TimeDependent else VectorFunction
-        return vec_func._new(3, 1, [comp1, comp2, comp3])
+        func = vec_func(self)
+        return func._new(3, 1, [comp1, comp2, comp3])
 
     def grad(self, shift=None):
         """
         Gradient of the VectorFunction, creates the gradient TensorFunction.
         """
-        func = tens_func(self, self)
-        comps = [[getattr(f, 'd%s' % d.name) for d in self.space_dimensions]
-                 for f in self]
+        func = tens_func(self)
+        ndim = len(self.space_dimensions)
+        shift_x0 = make_shift_x0(shift, (ndim, ndim))
+        comps = [[getattr(f, 'd%s' % d.name)(x0=shift_x0(shift, d, i, j))
+                  for j, d in enumerate(self.space_dimensions)]
+                 for i, f in enumerate(self)]
         return func._new(comps)
 
     def outer(self, other):
@@ -352,7 +359,7 @@ class VectorFunction(TensorFunction):
         return func._new(comps)
 
     def new_from_mat(self, mat):
-        func = vec_func(self, self)
+        func = vec_func(self)
         return func._new(self.rows, 1, mat)
 
 
@@ -391,16 +398,14 @@ mat_time_dict = {(True, True): TensorTimeFunction, (True, False): VectorTimeFunc
                  (False, True): TensorFunction, (False, False): VectorFunction}
 
 
-def vec_func(func1, func2):
-    f1 = getattr(func1, 'is_TimeDependent', False)
-    f2 = getattr(func2, 'is_TimeDependent', False)
-    return VectorTimeFunction if f1 or f2 else VectorFunction
+def vec_func(*funcs):
+    return (VectorTimeFunction if any([getattr(f, 'is_TimeDependent', False)
+            for f in funcs]) else VectorFunction)
 
 
-def tens_func(func1, func2):
-    f1 = getattr(func1, 'is_TimeDependent', False)
-    f2 = getattr(func2, 'is_TimeDependent', False)
-    return TensorTimeFunction if f1 or f2 else TensorFunction
+def tens_func(*funcs):
+    return (TensorTimeFunction if any([getattr(f, 'is_TimeDependent', False)
+            for f in funcs]) else TensorFunction)
 
 
 def sympify_tensor(arg):
diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py
index f78a8d064..a0a3ebcbe 100644
--- a/tests/test_derivatives.py
+++ b/tests/test_derivatives.py
@@ -465,36 +465,26 @@ class TestFD(object):
         b = np.dot(g_deriv.data.reshape(-1), f.data.reshape(-1))
         assert np.isclose(1 - a/b, 0, atol=1e-5)
 
-    @pytest.mark.parametrize('shift', [None, .5, -.5])
-    @pytest.mark.parametrize('ndim', [2, 3])
+    @pytest.mark.parametrize('shift, ndim', [(None, 2), (.5, 2), (.5, 3),
+                                             ((.5, .5, .5), 3)])
     def test_shifted_div(self, shift, ndim):
         grid = Grid(tuple([11]*ndim))
         f = Function(name="f", grid=grid, space_order=4)
         df = div(f, shift=shift).evaluate
         ref = 0
-        for d in grid.dimensions:
-            x0 = None if shift is None else d + shift * d.spacing
-            ref += getattr(f, 'd%s' % d.name)(x0=x0)
-        assert df == ref.evaluate
-
-    @pytest.mark.parametrize('shift', [None, .5, -.5])
-    @pytest.mark.parametrize('ndim', [2, 3])
-    def test_shifted_div_of_vectorfunction(self, shift, ndim):
-        grid = Grid(tuple([11]*ndim))
-        f = Function(name="f", grid=grid, space_order=4)
-        df = div(grad(f), shift=shift).evaluate
-        ref = 0
         for i, d in enumerate(grid.dimensions):
-            x0 = None if shift is None else d + shift * d.spacing
-            ref += getattr(grad(f)[i], 'd%s' % d.name)(x0=x0)
+            x0 = (None if shift is None else d + shift[i] * d.spacing if
+                  type(shift) is tuple else d + shift * d.spacing)
+            ref += getattr(f, 'd%s' % d.name)(x0=x0)
         assert df == ref.evaluate
 
-    @pytest.mark.parametrize('shift', [None, .5, -.5])
-    @pytest.mark.parametrize('ndim', [2, 3])
+    @pytest.mark.parametrize('shift, ndim', [(None, 2), (.5, 2), (.5, 3),
+                                             ((.5, .5, .5), 3)])
     def test_shifted_grad(self, shift, ndim):
         grid = Grid(tuple([11]*ndim))
         f = Function(name="f", grid=grid, space_order=4)
         g = grad(f, shift=shift).evaluate
-        for d, gi in zip(grid.dimensions, g):
-            x0 = None if shift is None else d + shift * d.spacing
+        for i, (d, gi) in enumerate(zip(grid.dimensions, g)):
+            x0 = (None if shift is None else d + shift[i] * d.spacing if
+                  type(shift) is tuple else d + shift * d.spacing)
             assert gi == getattr(f, 'd%s' % d.name)(x0=x0).evaluate
diff --git a/tests/test_tensors.py b/tests/test_tensors.py
index 66d5f6161..e3b3599dc 100644
--- a/tests/test_tensors.py
+++ b/tests/test_tensors.py
@@ -4,7 +4,7 @@ import sympy
 import pytest
 
 from devito import VectorFunction, TensorFunction, VectorTimeFunction, TensorTimeFunction
-from devito import Grid, Function, TimeFunction, Dimension, Eq
+from devito import Grid, Function, TimeFunction, Dimension, Eq, div, grad
 from devito.types import NODE
 
 
@@ -248,3 +248,58 @@ def test_partial_devito_tens(func1):
     for i in range(3):
         for j in range(3):
             assert f3[i, j] == sum(f2[i, k] * f1[j, k] for k in range(3))
+
+
+@pytest.mark.parametrize('shift, ndim', [(None, 2), (.5, 2), (.5, 3),
+                                         (tuple([tuple([.5]*3)]*3), 3)])
+def test_shifted_grad_of_vector(shift, ndim):
+    grid = Grid(tuple([11]*ndim))
+    f = VectorFunction(name="f", grid=grid, space_order=4)
+    gf = grad(f, shift=shift).evaluate
+
+    ref = []
+    for i in range(len(grid.dimensions)):
+        for j, d in enumerate(grid.dimensions):
+            x0 = (None if shift is None else d + shift[i][j] * d.spacing if
+                  type(shift) is tuple else d + shift * d.spacing)
+            ge = getattr(f[i], 'd%s' % d.name)(x0=x0)
+            ref.append(ge.evaluate)
+
+    for i, d in enumerate(gf):
+        assert d == ref[i]
+
+
+@pytest.mark.parametrize('shift, ndim', [(None, 2), (.5, 2), (.5, 3), ((.5, .5, .5), 3)])
+def test_shifted_div_of_vector(shift, ndim):
+    grid = Grid(tuple([11]*ndim))
+    v = VectorFunction(name="f", grid=grid, space_order=4)
+    df = div(v, shift=shift).evaluate
+    ref = 0
+
+    for i, d in enumerate(grid.dimensions):
+        x0 = (None if shift is None else d + shift[i] * d.spacing if
+              type(shift) is tuple else d + shift * d.spacing)
+        ref += getattr(v[i], 'd%s' % d.name)(x0=x0)
+
+    assert df == ref.evaluate
+
+
+@pytest.mark.parametrize('shift, ndim', [(None, 2), (.5, 2), (.5, 3),
+                                         (tuple([tuple([.5]*3)]*3), 3)])
+def test_shifted_div_of_tensor(shift, ndim):
+    grid = Grid(tuple([11]*ndim))
+    f = TensorFunction(name="f", grid=grid, space_order=4)
+    df = div(f, shift=shift).evaluate
+
+    ref = []
+    for i, a in enumerate(grid.dimensions):
+        elems = []
+        for j, d in reversed(list(enumerate(grid.dimensions))):
+            x0 = (None if shift is None else d + shift[i][j] * d.spacing if
+                  type(shift) is tuple else d + shift * d.spacing)
+            ge = getattr(f[i, j], 'd%s' % d.name)(x0=x0)
+            elems.append(ge.evaluate)
+        ref.append(sum(elems))
+
+    for i, d in enumerate(df):
+        assert d == ref[i]