api: update throughout for staggered setup

Author: mloubout

devito/finite_differences/derivative.py

@@ -405,7 +405,7 @@ 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 self.expr.is_Function:
             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,7 @@ 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)
+        dimensions, indices = cls.__indices_setup__(*args, **kwargs)
 
         # If it's an alias or simply has a different name, ignore `function`.
         # These cases imply the construction of a new AbstractFunction off
@@ -743,7 +743,6 @@ def __new__(cls, *args, **kwargs):
         # when executing __init_finalize__
         newobj._name = name
         newobj._dimensions = dimensions
-        newobj._staggered = staggered
         newobj._shape = cls.__shape_setup__(**kwargs)
         newobj._dtype = cls.__dtype_setup__(**kwargs)
 
@@ -926,11 +925,6 @@ def indices(self):
         """The indices of the object."""
         return DimensionTuple(*self.args, getters=self.dimensions)
 
-    @property
-    def staggered(self):
-        """The staggered indices of the object."""
-        return DimensionTuple(*self._staggered, getters=self.dimensions)
-
     @property
     def indices_ref(self):
         """The reference indices of the object (indices at first creation)."""
@@ -1496,6 +1490,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

@@ -26,7 +26,7 @@
 from devito.types.args import ArgProvider
 from devito.types.caching import CacheManager
 from devito.types.basic import AbstractFunction, Size
-from devito.types.utils import Buffer, DimensionTuple, NODE, CELL, host_layer
+from devito.types.utils import Buffer, DimensionTuple, NODE, CELL, host_layer, Staggering
 
 __all__ = ['Function', 'TimeFunction', 'SubFunction', 'TempFunction']
 
@@ -1010,6 +1010,10 @@ def _cache_meta(self):
     def __init_finalize__(self, *args, **kwargs):
         super().__init_finalize__(*args, **kwargs)
 
+        # Staggering
+        self._staggered = self.__staggered_setup__(self.dimensions,
+                                                   staggered=kwargs.get('staggered'))
+
         # Space order
         space_order = kwargs.get('space_order', 1)
         if isinstance(space_order, int):
@@ -1042,7 +1046,7 @@ def __fd_setup__(self):
 
     @cached_property
     def _fd_priority(self):
-        return 1 if self.staggered in [NODE, None] else 2
+        return 1 if self.staggered.on_node else 2
 
     @property
     def is_parameter(self):
@@ -1059,26 +1063,33 @@ def _eval_at(self, func):
         return self
 
     @classmethod
-    def __staggered_setup__(cls, dimensions, **kwargs):
+    def __staggered_setup__(cls, dimensions, staggered=None, **kwargs):
         """
         Setup staggering-related metadata. This method assigns:
 
             * 0 to non-staggered dimensions;
             * 1 to staggered dimensions.
         """
-        stagg = kwargs.get('staggered', None)
-        if stagg is CELL:
-            staggered = (sympy.S.One for d in dimensions)
-        elif stagg in [None, NODE]:
-            staggered = (sympy.S.Zero for d in dimensions)
-        elif all(is_integer(s) for s in as_tuple(stagg)):
+        if not staggered:
+            processed = ()
+        elif staggered is CELL:
+            processed = (sympy.S.One,)*len(dimensions)
+        elif staggered is NODE:
+            processed = (sympy.S.Zero,)*len(dimensions)
+        elif all(is_integer(s) for s in as_tuple(staggered)):
             # Staggering is already a tuple likely from rebuild
-            assert len(stagg) == len(dimensions)
-            return tuple(stagg)
+            assert len(staggered) == len(dimensions)
+            processed = staggered
         else:
-            staggered = (sympy.S.One if d in as_tuple(stagg) else sympy.S.Zero
-                         for d in dimensions)
-        return tuple(staggered)
+            processed = []
+            for d in dimensions:
+                if d in as_tuple(staggered):
+                    processed.append(sympy.S.One)
+                elif -d in as_tuple(staggered):
+                    processed.append(sympy.S.NegativeOne)
+                else:
+                    processed.append(sympy.S.Zero)
+        return tuple(processed)
 
     @classmethod
     def __indices_setup__(cls, *args, **kwargs):
@@ -1097,14 +1108,27 @@ 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))
-        return tuple(dimensions), tuple(staggered_indices), 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)
+
+    @property
+    def staggered(self):
+        """The staggered indices of the object."""
+        if self._staggered:
+            return Staggering(*self._staggered, getters=self.dimensions)
+        else:
+            return Staggering(getters=self.dimensions)
 
     @property
     def is_Staggered(self):
-        return self.staggered is not None
+        if not self.staggered:
+            return False
+        return True
 
     @classmethod
     def __shape_setup__(cls, **kwargs):
@@ -1392,7 +1416,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 +1430,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 +1469,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 self.staggered.on_node else 2.2
 
     @property
     def time_order(self):
@@ -1600,7 +1623,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,12 +69,13 @@ 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')
         dimensions = kwargs.get('dimensions')
-        inds, _, _ = Function.__indices_setup__(grid=grid,
-                                                dimensions=dimensions)
+        inds, _ = Function.__indices_setup__(grid=grid, dimensions=dimensions)
         self._space_dimensions = inds
 
     @classmethod

devito/types/utils.py

@@ -1,4 +1,5 @@
 from ctypes import POINTER, Structure
+from functools import cached_property
 
 from devito.tools import EnrichedTuple, Tag
 # Additional Function-related APIs
@@ -31,6 +32,13 @@ def __getitem_hook__(self, dim):
         raise KeyError
 
 
+class Staggering(DimensionTuple):
+
+    @cached_property
+    def on_node(self):
+        return not self or all(s == 0 for s in self)
+
+
 class IgnoreDimSort(tuple):
     """A tuple subclass used to wrap the implicit_dims to indicate
     that the topological sort of other dimensions should not occur."""