api: fix corner cases for x0/dims derivative specification

Author: mloubout

devito/finite_differences/derivative.py

@@ -8,7 +8,7 @@
 from .differentiable import Differentiable
 from .tools import direct, transpose
 from .rsfd import d45
-from devito.tools import as_mapper, as_tuple, filter_ordered, frozendict
+from devito.tools import as_mapper, as_tuple, filter_ordered, frozendict, is_integer
 from devito.types.utils import DimensionTuple
 
 __all__ = ['Derivative']
@@ -121,7 +121,8 @@ def __new__(cls, expr, *dims, **kwargs):
                     processed.append(i)
         obj._ppsubs = tuple(processed)
 
-        obj._x0 = frozendict(kwargs.get('x0', {}))
+        obj._x0 = cls._process_x0(obj._dims, **kwargs)
+
         return obj
 
     @classmethod
@@ -161,6 +162,10 @@ def _process_kwargs(cls, expr, *dims, **kwargs):
                     new_dims = (dims[0],)
                 else:
                     new_dims = tuple([dims[0]]*orders)
+        elif len(dims) == 2 and not isinstance(dims[1], Iterable) and is_integer(dims[1]):
+            # special case of single dimension and order
+            new_dims = (dims[0],)
+            orders = dims[1]
         else:
             # Iterable of 2-tuple, e.g. ((x, 2), (y, 3))
             new_dims = []
@@ -171,16 +176,16 @@ def _process_kwargs(cls, expr, *dims, **kwargs):
                     new_dims.extend([d[0] for _ in range(max(1, d[1]))])
                     orders.append(d[1])
                 else:
-                    new_dims.extend([d for _ in range(o)])
+                    new_dims.extend([d for _ in range(max(1, o))])
                     orders.append(o)
             new_dims = as_tuple(new_dims)
             orders = as_tuple(orders)
 
         # Finite difference orders depending on input dimension (.dt or .dx)
         fd_orders = kwargs.get('fd_order', tuple([expr.time_order if
                                                   getattr(d, 'is_Time', False) else
-                                                  expr.space_order for d in dims]))
-        if len(dims) == 1 and isinstance(fd_orders, Iterable):
+                                                  expr.space_order for d in new_dims]))
+        if len(new_dims) == 1 and isinstance(fd_orders, Iterable):
             fd_orders = fd_orders[0]
 
         # SymPy expects the list of variable w.r.t. which we differentiate to be a list
@@ -190,26 +195,36 @@ def _process_kwargs(cls, expr, *dims, **kwargs):
                           for s in filter_ordered(new_dims)]
         return new_dims, orders, fd_orders, variable_count
 
+    @classmethod
+    def _process_x0(cls, dims, **kwargs):
+        try:
+            x0 = frozendict(kwargs.get('x0', {}))
+        except TypeError:
+            # Only given a value
+            assert len(dims) == 1
+            x0 = frozendict({dims[0]: kwargs.get('x0')})
+
+        return x0
+
     def __call__(self, x0=None, fd_order=None, side=None, method=None):
+        x0 = self._process_x0(self.dims, x0=x0)
+        _x0 = frozendict({**self.x0, **x0})
         if self.ndims == 1:
             fd_order = fd_order or self._fd_order
             side = side or self._side
             method = method or self._method
-            x0 = {self.dims[0]: x0} if x0 is not None else self.x0
-            return self._new_from_self(fd_order=fd_order, side=side, x0=x0,
+            return self._new_from_self(fd_order=fd_order, side=side, x0=_x0,
                                        method=method)
 
         if side is not None:
             raise TypeError("Side only supported for first order single"
                             "Dimension derivative such as `.dxl` or .dx(side=left)")
         # Cross derivative
-        _x0 = dict(self._x0)
         _fd_order = dict(self.fd_order._getters)
         try:
             _fd_order.update(fd_order or {})
             _fd_order = tuple(_fd_order.values())
             _fd_order = DimensionTuple(*_fd_order, getters=self.dims)
-            _x0.update(x0 or {})
         except AttributeError:
             raise TypeError("Multi-dimensional Derivative, input expected as a dict")
 

tests/test_derivatives.py

@@ -676,6 +676,29 @@ def test_zero_fd_interp(self):
         finterp = Derivative(f, (x, 0), x0={x: x+x.spacing/2}).evaluate
         assert simplify(finterp - expected) == 0
 
+    def test_deriv_spec(self):
+        grid = Grid((11, 11))
+        x, y = grid.dimensions
+        f = Function(name="f", grid=grid, space_order=4)
+
+        assert f.dx(x0=x + x.spacing) == f.dx(x0={x: x + x.spacing})
+        assert Derivative(f, x, 1) == Derivative(f, (x, 1))
+
+        x0xy = {x: x+x.spacing/2, y: y+y.spacing/2}
+        dxy0 = Derivative(f, (x, 0), (y, 0), x0=x0xy)
+        dxy02 = Derivative(f, x, y, deriv_order=(0, 0), x0=x0xy)
+        assert dxy0 == dxy02
+        assert dxy0.dims == (x, y)
+        assert dxy0.deriv_order == (0, 0)
+        assert dxy0.fd_order == (4, 4)
+        assert dxy0.x0 == x0xy
+
+        dxy0 = Derivative(f, (x, 0), (y, 0), x0={y: y+y.spacing/2})
+        dxy02 = Derivative(f, x, y, deriv_order=(0, 0), x0={x: x+x.spacing/2})
+        assert dxy0 != dxy02
+        assert dxy0.x0 == {y: y+y.spacing/2}
+        assert dxy02.x0 == {x: x+x.spacing/2}
+
 
 class TestTwoStageEvaluation(object):