api: update throughout for staggered setup

Author: mloubout

devito/finite_differences/derivative.py

@@ -405,7 +405,8 @@ def _eval_at(self, func):
             return self
         # For basic equation of the form f = Derivative(g, ...) we can just
         # compare staggering
-        if self.expr.staggered == func.staggered:
+        if self.expr.staggered == func.staggered and \
+           not (self.expr.is_Add or self.expr.is_Mul):
             return self
 
         x0 = func.indices_ref.getters

devito/finite_differences/rsfd.py

@@ -1,6 +1,5 @@
 from functools import wraps
 
-from devito.types import NODE
 from devito.types.dimension import StencilDimension
 from .differentiable import Weights, DiffDerivative
 from .tools import generate_indices, fd_weights_registry
@@ -101,12 +100,7 @@ def check_staggering(func):
     def wrapper(expr, dim, x0=None, expand=True):
         grid = expr.grid
         x0 = {k: v for k, v in x0.items() if k.is_Space}
-        if expr.staggered is NODE or expr.staggered is None:
-            cond = x0 == {} or x0 == all_staggered(grid) or x0 == grid_node(grid)
-        elif expr.staggered == grid.dimensions:
-            cond = x0 == {} or x0 == all_staggered(grid) or x0 == grid_node(grid)
-        else:
-            cond = False
+        cond = x0 == {} or x0 == all_staggered(grid) or x0 == grid_node(grid)
         if cond:
             return func(expr, dim, x0=x0, expand=expand)
         else:
@@ -117,7 +111,8 @@ def wrapper(expr, dim, x0=None, expand=True):
 @check_staggering
 def d45(expr, dim, x0=None, expand=True):
     """
-    RSFD approximation of the derivative of `expr` along `dim` at point `x0`.
+    Rotated staggered grid finite-differences (RSFD) discretization
+    of the derivative of `expr` along `dim` at point `x0`.
 
     Parameters
     ----------
@@ -132,7 +127,8 @@ def d45(expr, dim, x0=None, expand=True):
     """
     # Make sure the grid supports RSFD
     if expr.grid.dim not in [2, 3]:
-        raise ValueError('RSFD only supported in 2D and 3D')
+        raise ValueError('Rotated staggered grid finite-differences (RSFD)'
+                         ' only supported in 2D and 3D')
 
     # Diagonals weights
     w = dir_weights[(dim.name, expr.grid.dim)]

devito/types/basic.py

@@ -699,7 +699,12 @@ def __new__(cls, *args, **kwargs):
         args, kwargs = cls.__args_setup__(*args, **kwargs)
 
         # Extract the `indices`, as perhaps they're explicitly provided
-        dimensions, indices, staggered = cls.__indices_setup__(*args, **kwargs)
+        dim_args = cls.__indices_setup__(*args, **kwargs)
+        try:
+            dimensions, indices, staggered = dim_args
+        except ValueError:
+            dimensions, indices = dim_args
+            staggered = (sympy.S.Zero,)*len(dimensions)
 
         # If it's an alias or simply has a different name, ignore `function`.
         # These cases imply the construction of a new AbstractFunction off
@@ -929,7 +934,10 @@ def indices(self):
     @property
     def staggered(self):
         """The staggered indices of the object."""
-        return DimensionTuple(*self._staggered, getters=self.dimensions)
+        if self._staggered:
+            return DimensionTuple(*self._staggered, getters=self.dimensions)
+        else:
+            return None
 
     @property
     def indices_ref(self):
@@ -1496,6 +1504,10 @@ def name(self):
         return self.__class__.__name__
 
     def _rebuild(self, *args, **kwargs):
+        # Plain `func` call (row, col, comps)
+        if not kwargs.keys() & self.__rkwargs__:
+            assert len(args) == 3
+            return self._new(*args, **kwargs)
         # We need to rebuild the components with the new name then
         # rebuild the matrix
         newname = kwargs.pop('name', self.name)

devito/types/dense.py

@@ -1042,7 +1042,7 @@ def __fd_setup__(self):
 
     @cached_property
     def _fd_priority(self):
-        return 1 if self.staggered in [NODE, None] else 2
+        return 1 if not self.staggered or all(s == 0 for s in self.staggered) else 2
 
     @property
     def is_parameter(self):
@@ -1066,18 +1066,26 @@ def __staggered_setup__(cls, dimensions, **kwargs):
             * 0 to non-staggered dimensions;
             * 1 to staggered dimensions.
         """
-        stagg = kwargs.get('staggered', None)
-        if stagg is CELL:
+        stagg = kwargs.get('staggered')
+        if stagg is None:
+            return tuple()
+        elif stagg is CELL:
             staggered = (sympy.S.One for d in dimensions)
-        elif stagg in [None, NODE]:
+        elif stagg in [NODE]:
             staggered = (sympy.S.Zero for d in dimensions)
         elif all(is_integer(s) for s in as_tuple(stagg)):
             # Staggering is already a tuple likely from rebuild
             assert len(stagg) == len(dimensions)
-            return tuple(stagg)
+            staggered = stagg
         else:
-            staggered = (sympy.S.One if d in as_tuple(stagg) else sympy.S.Zero
-                         for d in dimensions)
+            staggered = []
+            for d in dimensions:
+                if d in as_tuple(stagg):
+                    staggered.append(sympy.S.One)
+                elif -d in as_tuple(stagg):
+                    staggered.append(sympy.S.NegativeOne)
+                else:
+                    staggered.append(sympy.S.Zero)
         return tuple(staggered)
 
     @classmethod
@@ -1097,9 +1105,12 @@ def __indices_setup__(cls, *args, **kwargs):
             assert len(args) == len(dimensions)
             staggered_indices = tuple(args)
         else:
-            # Staggered indices
-            staggered_indices = (d + i * d.spacing / 2
-                                for d, i in zip(dimensions, staggered))
+            if not staggered:
+                staggered_indices = (d for d in dimensions)
+            else:
+                # Staggered indices
+                staggered_indices = (d + i * d.spacing / 2
+                                     for d, i in zip(dimensions, staggered))
         return tuple(dimensions), tuple(staggered_indices), staggered
 
     @property
@@ -1392,7 +1403,6 @@ def __fd_setup__(self):
     @classmethod
     def __indices_setup__(cls, *args, **kwargs):
         dimensions = kwargs.get('dimensions')
-        staggered = kwargs.get('staggered')
 
         if dimensions is None:
             save = kwargs.get('save')
@@ -1407,7 +1417,7 @@ def __indices_setup__(cls, *args, **kwargs):
             dimensions.insert(cls._time_position, time_dim)
 
         return Function.__indices_setup__(
-            *args, dimensions=dimensions, staggered=staggered
+            *args, dimensions=dimensions, staggered=kwargs.get('staggered')
         )
 
     @classmethod
@@ -1446,7 +1456,7 @@ def __shape_setup__(cls, **kwargs):
 
     @cached_property
     def _fd_priority(self):
-        return 2.1 if self.staggered in [NODE, None] else 2.2
+        return 2.1 if not self.staggered or all(s == 0 for s in self.staggered) else 2.2
 
     @property
     def time_order(self):
@@ -1600,7 +1610,7 @@ def __indices_setup__(cls, **kwargs):
         # Sanity check
         assert not any(d.is_NonlinearDerived for d in dimensions)
 
-        return dimensions, dimensions, (sympy.S.Zero for _ in dimensions)
+        return dimensions, dimensions
 
     def __halo_setup__(self, **kwargs):
         pointer_dim = kwargs.get('pointer_dim')

devito/types/tensor.py

@@ -69,6 +69,8 @@ class TensorFunction(AbstractTensor):
     _class_priority = 10
     _op_priority = Differentiable._op_priority + 1.
 
+    __rkwargs__ = AbstractTensor.__rkwargs__ + ('dimensions', 'space_order')
+
     def __init_finalize__(self, *args, **kwargs):
         super().__init_finalize__(*args, **kwargs)
         grid = kwargs.get('grid')