api: revamp staggered internal for correct rebuilding

Author: mloubout

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 = cls.__indices_setup__(*args, **kwargs)
+        dimensions, indices, staggered = 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,6 +743,7 @@ 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)
 
@@ -925,6 +926,11 @@ 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)."""
@@ -1428,8 +1434,7 @@ def _new(cls, *args, **kwargs):
                 return sympy.ImmutableDenseMatrix(*args)
             # Initialized with constructed object
             newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
-                                     grid=grid, dimensions=dimensions,
-                                     name=kwargs['name'])
+                                     grid=grid, dimensions=dimensions)
         else:
             # Initialize components and create new Matrix from standard
             # Devito inputs
@@ -1480,7 +1485,15 @@ def grid(self):
 
     @property
     def name(self):
-        return self._name
+        for c in self.values():
+            try:
+                return c.name.split('_')[0]
+            except AttributeError:
+                # `c` is not a devito object
+                pass
+        # If we end up here, then we have no devito objects
+        # in the matrix, so we ust return the class name
+        return self.__class__.__name__
 
     def _rebuild(self, *args, **kwargs):
         # We need to rebuild the components with the new name then
@@ -1489,7 +1502,7 @@ def _rebuild(self, *args, **kwargs):
         comps = [f.func(*args, name=f.name.replace(self.name, newname), **kwargs)
                  for f in self.flat()]
         # Rebuild the matrix with the new components
-        return self._new(comps, name=newname)
+        return self._new(comps)
 
     func = _rebuild
 

devito/types/dense.py

@@ -66,9 +66,6 @@ class DiscreteFunction(AbstractFunction, ArgProvider, Differentiable):
     __rkwargs__ = AbstractFunction.__rkwargs__ + ('staggered', 'coefficients')
 
     def __init_finalize__(self, *args, function=None, **kwargs):
-        # Staggering metadata
-        self._staggered = self.__staggered_setup__(**kwargs)
-
         # Now that *all* __X_setup__ hooks have been called, we can let the
         # superclass constructor do its job
         super().__init_finalize__(*args, **kwargs)
@@ -180,18 +177,6 @@ def __coefficients_setup__(self, **kwargs):
                                  " not %s" % (str(fd_weights_registry), coeffs))
         return coeffs
 
-    def __staggered_setup__(self, **kwargs):
-        """
-        Setup staggering-related metadata. This method assigns:
-
-            * 0 to non-staggered dimensions;
-            * 1 to staggered dimensions.
-        """
-        staggered = kwargs.get('staggered', None)
-        if staggered is CELL:
-            staggered = self.dimensions
-        return staggered
-
     @cached_property
     def _functions(self):
         return {self.function}
@@ -208,10 +193,6 @@ def _mem_external(self):
     def _mem_heap(self):
         return True
 
-    @property
-    def staggered(self):
-        return self._staggered
-
     @property
     def coefficients(self):
         """Form of the coefficients of the function."""
@@ -1077,34 +1058,49 @@ def _eval_at(self, func):
             return self.subs(mapper)
         return self
 
+    @classmethod
+    def __staggered_setup__(cls, dimensions, **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)):
+            # Staggering is already a tuple likely from rebuild
+            assert len(stagg) == len(dimensions)
+            return tuple(stagg)
+        else:
+            staggered = (sympy.S.One if d in as_tuple(stagg) else sympy.S.Zero
+                         for d in dimensions)
+        return tuple(staggered)
+
     @classmethod
     def __indices_setup__(cls, *args, **kwargs):
         grid = kwargs.get('grid')
         dimensions = kwargs.get('dimensions')
+        staggered = kwargs.get('staggered')
+
         if grid is None:
             if dimensions is None:
                 raise TypeError("Need either `grid` or `dimensions`")
         elif dimensions is None:
             dimensions = grid.dimensions
 
+        staggered = cls.__staggered_setup__(dimensions, staggered=staggered)
         if args:
             assert len(args) == len(dimensions)
-            return tuple(dimensions), tuple(args)
-
-        # Staggered indices
-        staggered = kwargs.get("staggered", None)
-        if staggered in [None, NODE]:
-            staggered_indices = dimensions
-        elif staggered == CELL:
-            staggered_indices = [d + d.spacing / 2 for d in dimensions]
+            staggered_indices = tuple(args)
         else:
-            mapper = {d: d for d in dimensions}
-            for s in as_tuple(staggered):
-                c, s = s.as_coeff_Mul()
-                mapper.update({s: s + c * s.spacing / 2})
-            staggered_indices = mapper.values()
-
-        return tuple(dimensions), tuple(staggered_indices)
+            # Staggered indices
+            staggered_indices = (d + i * d.spacing / 2
+                                for d, i in zip(dimensions, staggered))
+        return tuple(dimensions), tuple(staggered_indices), staggered
 
     @property
     def is_Staggered(self):
@@ -1604,7 +1600,7 @@ def __indices_setup__(cls, **kwargs):
         # Sanity check
         assert not any(d.is_NonlinearDerived for d in dimensions)
 
-        return dimensions, dimensions
+        return dimensions, dimensions, (sympy.S.Zero for _ in dimensions)
 
     def __halo_setup__(self, **kwargs):
         pointer_dim = kwargs.get('pointer_dim')

devito/types/tensor.py

@@ -73,8 +73,8 @@ 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

tests/test_staggered_utils.py

@@ -2,9 +2,9 @@
 import numpy as np
 from sympy import simplify
 
-from devito import (Function, Grid, NODE, VectorTimeFunction,
-                    TimeFunction, Eq, Operator, div)
-from devito.tools import powerset
+from devito import (Function, Grid, NODE, CELL, VectorTimeFunction,
+                    TimeFunction, Eq, Operator, div, Dimension)
+from devito.tools import powerset, as_tuple
 
 
 @pytest.mark.parametrize('ndim', [1, 2, 3])
@@ -160,3 +160,30 @@ def test_staggered_div():
     op2.apply(time_M=0)
 
     assert np.allclose(p1.data[:], p2.data[:], atol=0, rtol=1e-5)
+
+
+@pytest.mark.parametrize('stagg', [
+    'NODE', 'CELL', 'x', 'y', 'z',
+    '(x, y)', '(x, z)', '(y, z)', '(x, y, z)'])
+def test_staggered_rebuild(stagg):
+    grid = Grid(shape=(5, 5, 5))
+    x, y, z = grid.dimensions  # noqa
+    stagg = eval(stagg)
+
+    f = Function(name='f', grid=grid, space_order=4, staggered=stagg)
+    assert tuple(f.staggered.getters.keys()) == grid.dimensions
+
+    new_dims = (Dimension('x1'), Dimension('y1'), Dimension('z1'))
+    f2 = f.func(dimensions=new_dims)
+
+    assert f2.dimensions == new_dims
+    assert tuple(f2.staggered) == tuple(f.staggered)
+    assert tuple(f2.staggered.getters.keys()) == new_dims
+
+    # Check that rebuild correctly set the staggered indices
+    # with the new dimensions
+    for (d, nd) in zip(grid.dimensions, new_dims):
+        if d in as_tuple(stagg) or stagg is CELL:
+            assert f2.indices[nd] == nd + nd.spacing / 2
+        else:
+            assert f2.indices[nd] == nd

tests/test_tensors.py

@@ -455,11 +455,13 @@ def test_rebuild(func1):
         assert j.dimensions == i.dimensions
 
     new_dims = [Dimension(name=f'{i.name}1') for i in grid.dimensions]
+    if f1.is_TimeDependent:
+        new_dims = [f1[0].time_dim] + new_dims
     f3 = f1.func(dimensions=new_dims)
     assert f3.grid == grid
     assert f3.name == f1.name
 
     for (i, j) in zip(f1.flat(), f3.flat()):
         assert j.name == i.name
         assert j.grid == i.grid
-        assert j.dimensions == new_dims
+        assert j.dimensions == tuple(new_dims)