Merge pull request #2368 from devitocodes/zero-fd

api: cleanup FD tools and support zeroth order derivative

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
@@ -134,14 +135,11 @@ def _process_kwargs(cls, expr, *dims, **kwargs):
             fd_orders = kwargs.get('fd_order')
             deriv_orders = kwargs.get('deriv_order')
             if len(dims) == 1:
-                dims = tuple([dims[0]]*deriv_orders)
+                dims = tuple([dims[0]]*max(1, deriv_orders))
             variable_count = [sympy.Tuple(s, dims.count(s))
                               for s in filter_ordered(dims)]
             return dims, deriv_orders, fd_orders, variable_count
 
-        # Sanitise `dims`. ((x, 2), (y, 0)) is valid input, but (y, 0) should be dropped.
-        dims = tuple(d for d in dims if not (isinstance(d, Iterable) and d[1] == 0))
-
         # Check `dims`. It can be a single Dimension, an iterable of Dimensions, or even
         # an iterable of 2-tuple (Dimension, deriv_order)
         if len(dims) == 0:
@@ -154,62 +152,81 @@ def _process_kwargs(cls, expr, *dims, **kwargs):
                 orders = kwargs.get('deriv_order', dims[0][1])
                 if dims[0][1] != orders:
                     raise ValueError("Two different values of `deriv_order`")
-                new_dims = tuple([dims[0][0]]*dims[0][1])
+                new_dims = tuple([dims[0][0]]*max(1, dims[0][1]))
             else:
                 # Single Dimension
                 orders = kwargs.get('deriv_order', 1)
                 if isinstance(orders, Iterable):
                     orders = orders[0]
-                new_dims = tuple([dims[0]]*orders)
+                new_dims = tuple([dims[0]]*max(1, orders))
+        elif len(dims) == 2 and not isinstance(dims[1], Iterable) and is_integer(dims[1]):
+            # special case of single dimension and order
+            orders = dims[1]
+            new_dims = tuple([dims[0]]*max(1, orders))
         else:
             # Iterable of 2-tuple, e.g. ((x, 2), (y, 3))
             new_dims = []
             orders = []
             d_ord = kwargs.get('deriv_order', tuple([1]*len(dims)))
             for d, o in zip(dims, d_ord):
                 if isinstance(d, Iterable):
-                    new_dims.extend([d[0] for _ in range(d[1])])
+                    new_dims.extend([d[0]]*max(1, d[1]))
                     orders.append(d[1])
                 else:
-                    new_dims.extend([d for _ in range(o)])
+                    new_dims.extend([d]*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)
+        odims = filter_ordered(new_dims)
         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 odims]))
+        if len(odims) == 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
         # of 2-tuple `(s, count)` where s is the entity to diff wrt and count is the order
         # of the derivative
         variable_count = [sympy.Tuple(s, new_dims.count(s))
-                          for s in filter_ordered(new_dims)]
+                          for s in odims]
         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
+            _x0 = kwargs.get('x0')
+            assert len(dims) == 1 or _x0 is None
+            if _x0 is not None:
+                x0 = frozendict({dims[0]: _x0})
+            else:
+                x0 = frozendict({})
+
+        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")
 
@@ -340,7 +357,7 @@ def _eval_at(self, func):
         has to be computed at x=x + h_x/2.
         """
         # If an x0 already exists do not overwrite it
-        x0 = self.x0 or dict(func.indices_ref._getters)
+        x0 = self.x0 or func.indices_ref._getters
         if self.expr.is_Add:
             # If `expr` has both staggered and non-staggered terms such as
             # `(u(x + h_x/2) + v(x)).dx` then we exploit linearity of FD to split
@@ -412,7 +429,8 @@ def _eval_fd(self, expr, **kwargs):
                                    matvec=self.transpose, x0=self.x0, expand=expand)
         else:
             assert self.method == 'FD'
-            res = generic_derivative(expr, *self.dims, self.fd_order, self.deriv_order,
+            res = generic_derivative(expr, self.dims[0], as_tuple(self.fd_order)[0],
+                                     self.deriv_order,
                                      matvec=self.transpose, x0=self.x0, expand=expand)
 
         # Step 3: Apply substitutions

devito/finite_differences/differentiable.py

@@ -258,6 +258,7 @@ def __eq__(self, other):
         if ret is NotImplemented or not ret:
             # Non comparable or not equal as sympy objects
             return False
+
         return all(getattr(self, i, None) == getattr(other, i, None)
                    for i in self.__rkwargs__)
 
@@ -876,6 +877,11 @@ def __new__(cls, *args, base=None, **kwargs):
 
     func = DifferentiableOp._rebuild
 
+    # Since obj.base = base, then Differentiable.__eq__ leads to infinite recursion
+    # as it checks obj.base == other.base
+    __eq__ = sympy.Add.__eq__
+    __hash__ = sympy.Add.__hash__
+
     def _new_rawargs(self, *args, **kwargs):
         kwargs.pop('is_commutative', None)
         return self.func(*args, **kwargs)

devito/finite_differences/finite_difference.py

@@ -81,8 +81,14 @@ def first_derivative(expr, dim, fd_order=None, side=centered, matvec=direct, x0=
 
     Finally the x0 argument allows to choose the origin of the finite-difference
 
+    >>> first_derivative(f, dim=x, x0={x: x + x.spacing})
+    -f(x + h_x, y)/h_x + f(x + 2*h_x, y)/h_x
+
+    or specifying a specific location
+
     >>> first_derivative(f, dim=x, x0={x: 1})
-    -f(1, y)/h_x + f(h_x + 1, y)/h_x
+    f(1, y)/h_x - f(1 - h_x, y)/h_x
+
     """
     fd_order = fd_order or expr.space_order
     deriv_order = 1
@@ -155,9 +161,11 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, **kwargs):
     Finally the x0 argument allows to choose the origin of the finite-difference
 
     >>> cross_derivative(f*g, dims=(x, y), fd_order=(2, 2), deriv_order=(1, 1), \
-    x0={x: 1, y: 2})
-    (-1/h_y)*(-f(1, 2)*g(1, 2)/h_x + f(h_x + 1, 2)*g(h_x + 1, 2)/h_x) + (-f(1, h_y + 2)*\
-g(1, h_y + 2)/h_x + f(h_x + 1, h_y + 2)*g(h_x + 1, h_y + 2)/h_x)/h_y
+    x0={x: x + x.spacing, y: y + y.spacing})
+    (-1/h_y)*(-f(x + h_x, y + h_y)*g(x + h_x, y + h_y)/h_x + \
+f(x + 2*h_x, y + h_y)*g(x + 2*h_x, y + h_y)/h_x) + \
+(-f(x + h_x, y + 2*h_y)*g(x + h_x, y + 2*h_y)/h_x + \
+f(x + 2*h_x, y + 2*h_y)*g(x + 2*h_x, y + 2*h_y)/h_x)/h_y
     """
     x0 = x0 or {}
     for d, fd, dim in zip(deriv_order, fd_order, dims):
@@ -208,6 +216,10 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
     if deriv_order == 1 and fd_order == 2 and coefficients != 'symbolic':
         fd_order = 1
 
+    # Zeroth order derivative is just the expression itself if not shifted
+    if deriv_order == 0 and not x0:
+        return expr
+
     # Enforce stable time coefficients
     if dim.is_Time and coefficients != 'symbolic':
         coefficients = 'taylor'

devito/finite_differences/rsfd.py

@@ -3,7 +3,7 @@
 from devito.types import NODE
 from devito.types.dimension import StencilDimension
 from .differentiable import Weights, DiffDerivative
-from .tools import generate_indices_staggered, fd_weights_registry
+from .tools import generate_indices, fd_weights_registry
 
 __all__ = ['drot', 'd45']
 
@@ -42,7 +42,7 @@ def drot(expr, dim, dir=1, x0=None):
     s = 2**(expr.grid.dim - 1)
 
     # Center point and indices
-    start, indices = generate_indices_staggered(expr, dim, expr.space_order, x0=x0)
+    indices, start = generate_indices(expr, dim, expr.space_order, x0=x0)
 
     # a-dimensional center for FD coefficients.
     adim_start = x0.get(dim, expr.indices_ref[dim]).subs({dim: 0, dim.spacing: 1})

devito/finite_differences/tools.py

@@ -275,128 +275,35 @@ def generate_indices(expr, dim, order, side=None, matvec=None, x0=None):
     -------
     An IndexSet, representing an ordered list of indices.
     """
-    if expr.is_Staggered and not dim.is_Time and side is None:
-        x0, indices = generate_indices_staggered(expr, dim, order, side=side, x0=x0)
-    else:
-        x0 = (x0 or {dim: dim}).get(dim, dim)
-        # Check if called from first_derivative()
-        indices = generate_indices_cartesian(expr, dim, order, side, x0)
-
-    assert isinstance(indices, IndexSet)
-
-    return indices, x0
-
-
-def generate_indices_cartesian(expr, dim, order, side, x0):
-    """
-    Indices for the finite-difference scheme on a cartesian grid.
-
-    Parameters
-    ----------
-    expr : expr-like
-        Expression that is differentiated.
-    dim : Dimension
-        Dimensions w.r.t which the derivative is taken.
-    order : int
-        Order of the finite-difference scheme.
-    side : Side
-        Side of the scheme (centered, left, right).
-    x0 : dict of {Dimension: Dimension or expr-like or Number}
-        Origin of the scheme, ie. `x`, `x + .5 * x.spacing`, ...
-
-    Returns
-    -------
-    An IndexSet, representing an ordered list of indices.
-    """
-    shift = 0
-    # Shift if `x0` is not on the grid
-    offset_c = 0 if sympify(x0).is_Integer else (dim - x0)/dim.spacing
-    offset_c = np.sign(offset_c) * (offset_c % 1)
-    offset = offset_c * dim.spacing
-    # Spacing
-    diff = dim.spacing
-    if side in [left, right]:
-        shift = 1
-        diff *= side.val
-    # Indices
-    if order < 2:
-        indices = [x0, x0 + diff] if offset == 0 else [x0 - offset, x0 + offset]
-        return IndexSet(dim, indices)
-    else:
-        # Left and right max offsets for indices
-        o_min = -order//2 + int(np.ceil(-offset_c))
-        o_max = order//2 - int(np.ceil(offset_c))
-
-        d = make_stencil_dimension(expr, o_min, o_max)
-        iexpr = x0 + (d + shift) * diff + offset
-        return IndexSet(dim, expr=iexpr)
-
-
-def generate_indices_staggered(expr, dim, order, side=None, x0=None):
-    """
-    Indices for the finite-difference scheme on a staggered grid.
-
-    Parameters
-    ----------
-    expr : expr-like
-        Expression that is differentiated.
-    dim : Dimension
-        Dimensions w.r.t which the derivative is taken.
-    order : int
-        Order of the finite-difference scheme.
-    side : Side, optional
-        Side of the scheme (centered, left, right).
-    x0 : dict of {Dimension: Dimension or expr-like or Number}, optional
-        Origin of the scheme, ie. `x`, `x + .5 * x.spacing`, ...
-
-    Returns
-    -------
-    An IndexSet, representing an ordered list of indices.
-    """
-    diff = dim.spacing
-    start = (x0 or {}).get(dim) or expr.indices_ref[dim]
-    try:
-        ind0 = expr.indices_ref[dim]
-    except AttributeError:
-        ind0 = start
-
-    if start != ind0:
-        if order < 2:
-            indices = [start - diff/2, start + diff/2]
-            indices = IndexSet(dim, indices)
+    # Evaluation point
+    x0 = sympify(((x0 or {}).get(dim) or expr.indices_ref[dim]))
+
+    # If provided a pure number, assume it's a valid index
+    if x0.is_Number:
+        d = make_stencil_dimension(expr, -order//2, order//2)
+        iexpr = x0 + d * dim.spacing
+        return IndexSet(dim, expr=iexpr), x0
+
+    # Shift for side
+    side = side or centered
+
+    # Evaluation point relative to the expression's grid
+    mid = (x0 - expr.indices_ref[dim]).subs({dim: 0, dim.spacing: 1})
+
+    # Indices range
+    o_min = int(np.ceil(mid - order/2)) + side.val
+    o_max = int(np.floor(mid + order/2)) + side.val
+    if o_max == o_min:
+        if dim.is_Time or not expr.is_Staggered:
+            o_max += 1
         else:
-            o_min = -order//2 + 1
-            o_max = order//2
-
-            d = make_stencil_dimension(expr, o_min, o_max)
-            iexpr = start - diff/2 + d * diff
-            indices = IndexSet(dim, expr=iexpr)
-    else:
-        if order < 2:
-            indices = [start, start - diff]
-            indices = IndexSet(dim, indices)
-        else:
-            if x0 is None or order % 2 == 0:
-                # No _eval_at or even order derivatives
-                # keep the centered indices
-                o_min = -order//2
-                o_max = order//2
-            elif start is dim:
-                # Staggered FD requires half cell shift
-                # for stability
-                o_min = -order//2 + 1
-                o_max = order//2
-                start = dim + diff/2
-            else:
-                o_min = -order//2
-                o_max = order//2 - 1
-                start = dim
+            o_min -= 1
 
-            d = make_stencil_dimension(expr, o_min, o_max)
-            iexpr = ind0 + d * diff
-            indices = IndexSet(dim, expr=iexpr)
+    # StencilDimension and expression
+    d = make_stencil_dimension(expr, o_min, o_max)
+    iexpr = expr.indices_ref[dim] + d * dim.spacing
 
-    return start, indices
+    return IndexSet(dim, expr=iexpr), x0
 
 
 def make_shift_x0(shift, ndim):