api: fix staggering handling at rebuild

Author: mloubout

devito/types/dense.py

@@ -1143,6 +1143,9 @@ def __staggered_setup__(cls, dimensions, staggered=None, **kwargs):
             * 0 to non-staggered dimensions;
             * 1 to staggered dimensions.
         """
+        if isinstance(staggered, Staggering):
+            staggered = staggered._ref
+
         if not staggered:
             processed = ()
         elif staggered is CELL:

devito/types/tensor.py

@@ -20,6 +20,18 @@
 __all__ = ['TensorFunction', 'TensorTimeFunction', 'VectorFunction', 'VectorTimeFunction']
 
 
+def staggering(stagg, i, j, d, dims):
+    if stagg is None:
+        # No input
+        return NODE if i == j else (d, dims[j])
+    elif isinstance(stagg, (tuple, list)):
+        # User input as list or tuple
+        return stagg[i][j]
+    elif isinstance(stagg, AbstractTensor):
+        # From rebuild/tensor property. Indexed as a sympy Matrix
+        return stagg[i, j]
+
+
 class TensorFunction(AbstractTensor):
     """
     Tensor valued Function represented as a Matrix.
@@ -128,8 +140,7 @@ def __subfunc_setup__(cls, *args, **kwargs):
             start = i if (symm or diag) else 0
             stop = i + 1 if diag else len(dims)
             for j in range(start, stop):
-                staggj = (stagg[i][j] if stagg is not None
-                          else (NODE if i == j else (d, dims[j])))
+                staggj = staggering(stagg, i, j, d, dims)
                 sub_kwargs = cls._component_kwargs((i, j), **kwargs)
                 sub_kwargs.update({'name': f"{name}_{d.name}{dims[j].name}",
                                    'staggered': staggj})

devito/types/utils.py

@@ -57,6 +57,13 @@ class Staggering(DimensionTuple):
     def on_node(self):
         return not self or all(s == 0 for s in self)
 
+    @property
+    def _ref(self):
+        if self.on_node:
+            return NODE
+        else:
+            return tuple(d for d, s in zip(self.getters, self, strict=True) if s == 1)
+
 
 class IgnoreDimSort(tuple):
     """A tuple subclass used to wrap the implicit_dims to indicate