From c1c58062c1776f2efae407710ae87b002b2b889c Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Fri, 11 Apr 2025 13:53:20 -0400
Subject: [PATCH 1/4] api: fix tensor rebuild

---
 devito/types/basic.py  | 398 +++++++++++++++++++++--------------------
 devito/types/dense.py  |   4 +-
 devito/types/tensor.py |   1 +
 tests/test_tensors.py  |  25 +++
 4 files changed, 237 insertions(+), 191 deletions(-)

diff --git a/devito/types/basic.py b/devito/types/basic.py
index 04e4281f12..03a95fd12c 100644
--- a/devito/types/basic.py
+++ b/devito/types/basic.py
@@ -620,196 +620,6 @@ def _arg_values(self, **kwargs):
             return self._arg_defaults(**kwargs)
 
 
-class AbstractTensor(sympy.ImmutableDenseMatrix, Basic, Pickable, Evaluable):
-
-    """
-    Base class for vector and tensor valued functions. It inherits from and
-    mimicks the behavior of a sympy.ImmutableDenseMatrix.
-
-
-    The sub-hierachy is as follows
-
-                         AbstractTensor
-                                |
-                          TensorFunction
-                                |
-                 ---------------------------------
-                 |                               |
-          VectorFunction                 TensorTimeFunction
-                        \\-------\\              |
-                                 \\------- VectorTimeFunction
-
-    There are four relevant AbstractTensor sub-types: ::
-
-        * TensorFunction: A space-varying tensor valued function.
-        * VectorFunction: A space-varying vector valued function.
-        * TensorTimeFunction: A time-space-varying tensor valued function.
-        * VectorTimeFunction: A time-space-varying vector valued function.
-    """
-
-    # SymPy attributes
-    is_MatrixLike = True
-    is_Matrix = True
-
-    # Devito attributes
-    is_AbstractTensor = True
-    is_TensorValued = True
-    is_VectorValued = False
-
-    @classmethod
-    def _new(cls, *args, **kwargs):
-        if args:
-            try:
-                # Constructor if input is (rows, cols, lambda)
-                newobj = super()._new(*args)
-            except ValueError:
-                # Constructor if input is list of list as (row, cols, list_of_list)
-                # doesn't work as it expects a flattened.
-                newobj = super()._new(args[2])
-
-            # Filter grid and dimensions
-            grid, dimensions = newobj._infer_dims()
-            if grid is None and dimensions is None:
-                return sympy.ImmutableDenseMatrix(*args)
-            # Initialized with constructed object
-            newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
-                                     grid=grid, dimensions=dimensions)
-        else:
-            # Initialize components and create new Matrix from standard
-            # Devito inputs
-            comps = cls.__subfunc_setup__(*args, **kwargs)
-            newobj = super()._new(comps)
-            newobj.__init_finalize__(*args, **kwargs)
-
-        return newobj
-
-    @classmethod
-    def _fromrep(cls, rep):
-        """
-        This the new constructor mechanism for matrices in sympy 1.9.
-        Standard new object go through `_new` but arithmetic operations directly use
-        the representation based one.
-        This class method is only accessible from an existing AbstractTensor
-        that contains a grid or dimensions.
-        """
-        newobj = super()._fromrep(rep)
-        grid, dimensions = newobj._infer_dims()
-        try:
-            # This is needed when `_fromrep` is called directly in 1.9
-            # for example with mul.
-            newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
-                                     grid=grid, dimensions=dimensions)
-        except TypeError:
-            # We can end up here when `_fromrep` is called through the default _new
-            # when input `comps` don't have grid or dimensions. For example
-            # `test_non_devito_tens` in `test_tensor.py`.
-            pass
-        return newobj
-
-    @classmethod
-    def __subfunc_setup__(cls, *args, **kwargs):
-        """Setup each component of the tensor as a Devito type."""
-        return []
-
-    @property
-    def grid(self):
-        """
-        A Tensor is expected to have all its components defined over the same grid
-        """
-        grids = {getattr(c, 'grid', None) for c in self.flat()} - {None}
-        if len(grids) == 0:
-            return None
-        assert len(grids) == 1
-        return grids.pop()
-
-    def _infer_dims(self):
-        grids = {getattr(c, 'grid', None) for c in self.flat()} - {None}
-        grids = {g.root for g in grids}
-        dimensions = {d for c in self.flat()
-                      for d in getattr(c, 'dimensions', ())} - {None}
-        # If none of the components are devito objects, returns a sympy Matrix
-        if len(grids) == 0 and len(dimensions) == 0:
-            return None, None
-        elif len(grids) > 0:
-            dimensions = None
-            assert len(grids) == 1
-            grid = grids.pop()
-        else:
-            grid = None
-            dimensions = tuple(dimensions)
-
-        return grid, dimensions
-
-    def flat(self):
-        try:
-            return super().flat()
-        except AttributeError:
-            return self._mat
-
-    def __init_finalize__(self, *args, **kwargs):
-        pass
-
-    __hash__ = sympy.ImmutableDenseMatrix.__hash__
-
-    def doit(self, **hint):
-        return self
-
-    def transpose(self, inner=True):
-        new = super().transpose()
-        if inner:
-            return new.applyfunc(lambda x: getattr(x, 'T', x))
-        return new
-
-    def adjoint(self, inner=True):
-        # Real valued adjoint is transpose
-        return self.transpose(inner=inner)
-
-    @call_highest_priority('__radd__')
-    def __add__(self, other):
-        try:
-            # Most case support sympy add
-            tsum = super().__add__(other)
-        except TypeError:
-            # Sympy doesn't support add with scalars
-            tsum = self.applyfunc(lambda x: x + other)
-
-        # As of sympy 1.13, super does not throw an exception but
-        # only returns NotImplemented for some internal dispatch.
-        if tsum is NotImplemented:
-            return self.applyfunc(lambda x: x + other)
-
-        return tsum
-
-    def _eval_matrix_mul(self, other):
-        """
-        Copy paste from sympy to avoid explicit call to sympy.Add
-        TODO: fix inside sympy
-        """
-        other_len = other.rows*other.cols
-        new_len = self.rows*other.cols
-        new_mat = [self.zero]*new_len
-
-        # If we multiply an n x 0 with a 0 x m, the
-        # expected behavior is to produce an n x m matrix of zeros
-        if self.cols != 0 and other.rows != 0:
-            self_cols = self.cols
-            mat = self.flat()
-            try:
-                other_mat = other.flat()
-            except AttributeError:
-                other_mat = other._mat
-            for i in range(new_len):
-                row, col = i // other.cols, i % other.cols
-                row_indices = range(self_cols*row, self_cols*(row+1))
-                col_indices = range(col, other_len, other.cols)
-                vec = [mat[a]*other_mat[b] for a, b in zip(row_indices, col_indices)]
-                new_mat[i] = sum(vec)
-
-        # Get new class and return product
-        newcls = self.classof_prod(other, other.cols)
-        return newcls._new(self.rows, other.cols, new_mat, copy=False)
-
-
 class AbstractFunction(sympy.Function, Basic, Pickable, Evaluable):
 
     """
@@ -1563,6 +1373,214 @@ def __getnewargs_ex__(self):
         return args, kwargs
 
 
+class AbstractTensor(sympy.ImmutableDenseMatrix, Basic, Pickable, Evaluable):
+
+    """
+    Base class for vector and tensor valued functions. It inherits from and
+    mimicks the behavior of a sympy.ImmutableDenseMatrix.
+
+
+    The sub-hierachy is as follows
+
+                         AbstractTensor
+                                |
+                          TensorFunction
+                                |
+                 ---------------------------------
+                 |                               |
+          VectorFunction                 TensorTimeFunction
+                        \\-------\\              |
+                                 \\------- VectorTimeFunction
+
+    There are four relevant AbstractTensor sub-types: ::
+
+        * TensorFunction: A space-varying tensor valued function.
+        * VectorFunction: A space-varying vector valued function.
+        * TensorTimeFunction: A time-space-varying tensor valued function.
+        * VectorTimeFunction: A time-space-varying vector valued function.
+    """
+
+    # SymPy attributes
+    is_MatrixLike = True
+    is_Matrix = True
+
+    # Devito attributes
+    is_AbstractTensor = True
+    is_TensorValued = True
+    is_VectorValued = False
+
+    __rkwargs__ = AbstractFunction.__rkwargs__
+
+    @classmethod
+    def _new(cls, *args, **kwargs):
+        if args:
+            try:
+                # Constructor if input is (rows, cols, lambda)
+                newobj = super()._new(*args)
+            except ValueError:
+                # Constructor if input is list of list as (row, cols, list_of_list)
+                # doesn't work as it expects a flattened.
+                newobj = super()._new(args[2])
+
+            # Filter grid and dimensions
+            grid, dimensions = newobj._infer_dims()
+            if grid is None and dimensions is None:
+                return sympy.ImmutableDenseMatrix(*args)
+            # Initialized with constructed object
+            newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
+                                     grid=grid, dimensions=dimensions,
+                                     name=kwargs['name'])
+        else:
+            # Initialize components and create new Matrix from standard
+            # Devito inputs
+            comps = cls.__subfunc_setup__(*args, **kwargs)
+            newobj = super()._new(comps)
+            newobj.__init_finalize__(*args, **kwargs)
+
+        return newobj
+
+    @classmethod
+    def _fromrep(cls, rep):
+        """
+        This the new constructor mechanism for matrices in sympy 1.9.
+        Standard new object go through `_new` but arithmetic operations directly use
+        the representation based one.
+        This class method is only accessible from an existing AbstractTensor
+        that contains a grid or dimensions.
+        """
+        newobj = super()._fromrep(rep)
+        grid, dimensions = newobj._infer_dims()
+        try:
+            # This is needed when `_fromrep` is called directly in 1.9
+            # for example with mul.
+            newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
+                                     grid=grid, dimensions=dimensions)
+        except TypeError:
+            # We can end up here when `_fromrep` is called through the default _new
+            # when input `comps` don't have grid or dimensions. For example
+            # `test_non_devito_tens` in `test_tensor.py`.
+            pass
+        return newobj
+
+    @classmethod
+    def __subfunc_setup__(cls, *args, **kwargs):
+        """Setup each component of the tensor as a Devito type."""
+        return []
+
+    @property
+    def grid(self):
+        """
+        A Tensor is expected to have all its components defined over the same grid
+        """
+        grids = {getattr(c, 'grid', None) for c in self.flat()} - {None}
+        if len(grids) == 0:
+            return None
+        assert len(grids) == 1
+        return grids.pop()
+
+    @property
+    def name(self):
+        return self._name
+
+    def _rebuild(self, *args, **kwargs):
+        # We need to rebuild the components with the new name then
+        # rebuild the matrix
+        newname = kwargs.pop('name', self.name)
+        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)
+
+    func = _rebuild
+
+    def _infer_dims(self):
+        grids = {getattr(c, 'grid', None) for c in self.flat()} - {None}
+        grids = {g.root for g in grids}
+        dimensions = {d for c in self.flat()
+                      for d in getattr(c, 'dimensions', ())} - {None}
+        # If none of the components are devito objects, returns a sympy Matrix
+        if len(grids) == 0 and len(dimensions) == 0:
+            return None, None
+        elif len(grids) > 0:
+            dimensions = None
+            assert len(grids) == 1
+            grid = grids.pop()
+        else:
+            grid = None
+            dimensions = tuple(dimensions)
+
+        return grid, dimensions
+
+    def flat(self):
+        try:
+            return super().flat()
+        except AttributeError:
+            return self._mat
+
+    def __init_finalize__(self, *args, **kwargs):
+        self._name = kwargs.get('name', None)
+
+    __hash__ = sympy.ImmutableDenseMatrix.__hash__
+
+    def doit(self, **hint):
+        return self
+
+    def transpose(self, inner=True):
+        new = super().transpose()
+        if inner:
+            return new.applyfunc(lambda x: getattr(x, 'T', x))
+        return new
+
+    def adjoint(self, inner=True):
+        # Real valued adjoint is transpose
+        return self.transpose(inner=inner)
+
+    @call_highest_priority('__radd__')
+    def __add__(self, other):
+        try:
+            # Most case support sympy add
+            tsum = super().__add__(other)
+        except TypeError:
+            # Sympy doesn't support add with scalars
+            tsum = self.applyfunc(lambda x: x + other)
+
+        # As of sympy 1.13, super does not throw an exception but
+        # only returns NotImplemented for some internal dispatch.
+        if tsum is NotImplemented:
+            return self.applyfunc(lambda x: x + other)
+
+        return tsum
+
+    def _eval_matrix_mul(self, other):
+        """
+        Copy paste from sympy to avoid explicit call to sympy.Add
+        TODO: fix inside sympy
+        """
+        other_len = other.rows*other.cols
+        new_len = self.rows*other.cols
+        new_mat = [self.zero]*new_len
+
+        # If we multiply an n x 0 with a 0 x m, the
+        # expected behavior is to produce an n x m matrix of zeros
+        if self.cols != 0 and other.rows != 0:
+            self_cols = self.cols
+            mat = self.flat()
+            try:
+                other_mat = other.flat()
+            except AttributeError:
+                other_mat = other._mat
+            for i in range(new_len):
+                row, col = i // other.cols, i % other.cols
+                row_indices = range(self_cols*row, self_cols*(row+1))
+                col_indices = range(col, other_len, other.cols)
+                vec = [mat[a]*other_mat[b] for a, b in zip(row_indices, col_indices)]
+                new_mat[i] = sum(vec)
+
+        # Get new class and return product
+        newcls = self.classof_prod(other, other.cols)
+        return newcls._new(self.rows, other.cols, new_mat, copy=False)
+
+
 # Extended SymPy hierarchy follows, for essentially two reasons:
 # - To keep track of `function`
 # - To override SymPy caching behaviour
diff --git a/devito/types/dense.py b/devito/types/dense.py
index 0b2ed9f828..3182421ada 100644
--- a/devito/types/dense.py
+++ b/devito/types/dense.py
@@ -1093,8 +1093,10 @@ def __indices_setup__(cls, *args, **kwargs):
 
         # Staggered indices
         staggered = kwargs.get("staggered", None)
-        if staggered in [CELL, NODE]:
+        if staggered in [None, NODE]:
             staggered_indices = dimensions
+        elif staggered == CELL:
+            staggered_indices = [d + d.spacing / 2 for d in dimensions]
         else:
             mapper = {d: d for d in dimensions}
             for s in as_tuple(staggered):
diff --git a/devito/types/tensor.py b/devito/types/tensor.py
index 2fdd4cc1ff..192d224364 100644
--- a/devito/types/tensor.py
+++ b/devito/types/tensor.py
@@ -70,6 +70,7 @@ class TensorFunction(AbstractTensor):
     _op_priority = Differentiable._op_priority + 1.
 
     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,
diff --git a/tests/test_tensors.py b/tests/test_tensors.py
index fdad398a36..58b19e081c 100644
--- a/tests/test_tensors.py
+++ b/tests/test_tensors.py
@@ -438,3 +438,28 @@ def test_custom_coeffs_tensor():
         derivs = retrieve_derivatives(dtau)
         for drv in derivs:
             assert list(drv.weights) == c
+
+
+@pytest.mark.parametrize('func1', [TensorFunction, TensorTimeFunction,
+                                   VectorFunction, VectorTimeFunction])
+def test_rebuild(func1):
+    grid = Grid(tuple([5]*3))
+    f1 = func1(name="f1", grid=grid)
+    f2 = f1.func(name="f2")
+    assert f1.grid == f2.grid
+    assert f2.name == 'f2'
+
+    for (i, j) in zip(f1.flat(), f2.flat()):
+        assert j.name == i.name.replace('f1', 'f2')
+        assert j.grid == i.grid
+        assert j.dimensions == i.dimensions
+
+    new_dims = [Dimension(name=f'{i.name}1') for i in grid.dimensions]
+    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

From 5dd4f97580d149607fb3785d58388d4fdcdc8cf8 Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Mon, 14 Apr 2025 09:44:31 -0400
Subject: [PATCH 2/4] api: revamp staggered internal for correct rebuilding

---
 devito/types/basic.py         | 23 +++++++++---
 devito/types/dense.py         | 66 ++++++++++++++++-------------------
 devito/types/tensor.py        |  4 +--
 tests/test_staggered_utils.py | 33 ++++++++++++++++--
 tests/test_tensors.py         |  4 ++-
 5 files changed, 84 insertions(+), 46 deletions(-)

diff --git a/devito/types/basic.py b/devito/types/basic.py
index 03a95fd12c..7e01f407c8 100644
--- a/devito/types/basic.py
+++ b/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
 
diff --git a/devito/types/dense.py b/devito/types/dense.py
index 3182421ada..5676826034 100644
--- a/devito/types/dense.py
+++ b/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')
diff --git a/devito/types/tensor.py b/devito/types/tensor.py
index 192d224364..c36702bea1 100644
--- a/devito/types/tensor.py
+++ b/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
diff --git a/tests/test_staggered_utils.py b/tests/test_staggered_utils.py
index ece859cd95..4728892f0c 100644
--- a/tests/test_staggered_utils.py
+++ b/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
diff --git a/tests/test_tensors.py b/tests/test_tensors.py
index 58b19e081c..268bd9ea5c 100644
--- a/tests/test_tensors.py
+++ b/tests/test_tensors.py
@@ -455,6 +455,8 @@ 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
@@ -462,4 +464,4 @@ def test_rebuild(func1):
     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)

From 29fc0c96df361a15dd2e910f6069b473f5b44694 Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Mon, 14 Apr 2025 12:27:22 -0400
Subject: [PATCH 3/4] api: update throughout for staggered setup

---
 devito/finite_differences/derivative.py |  2 +-
 devito/finite_differences/rsfd.py       | 14 ++---
 devito/types/basic.py                   | 12 ++---
 devito/types/dense.py                   | 69 ++++++++++++++++---------
 devito/types/tensor.py                  |  5 +-
 devito/types/utils.py                   |  8 +++
 6 files changed, 68 insertions(+), 42 deletions(-)

diff --git a/devito/finite_differences/derivative.py b/devito/finite_differences/derivative.py
index c73c1071de..8bc36a7e07 100644
--- a/devito/finite_differences/derivative.py
+++ b/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
diff --git a/devito/finite_differences/rsfd.py b/devito/finite_differences/rsfd.py
index 1d77c26e05..6f6ecf8a93 100644
--- a/devito/finite_differences/rsfd.py
+++ b/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)]
diff --git a/devito/types/basic.py b/devito/types/basic.py
index 7e01f407c8..593e3ffff7 100644
--- a/devito/types/basic.py
+++ b/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)
diff --git a/devito/types/dense.py b/devito/types/dense.py
index 5676826034..d763ea5b71 100644
--- a/devito/types/dense.py
+++ b/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')
diff --git a/devito/types/tensor.py b/devito/types/tensor.py
index c36702bea1..54cf280880 100644
--- a/devito/types/tensor.py
+++ b/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
diff --git a/devito/types/utils.py b/devito/types/utils.py
index 7673489149..75c5885b94 100644
--- a/devito/types/utils.py
+++ b/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."""

From 9c9fe55afeb3d760debe4a11262a2576594fc91c Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Wed, 16 Apr 2025 22:21:51 -0400
Subject: [PATCH 4/4] api: add support for diag(tensor) and diag(vector)

---
 devito/finite_differences/operators.py | 10 +++++++--
 devito/types/basic.py                  |  4 +++-
 devito/types/dense.py                  | 14 ++++--------
 tests/test_tensors.py                  | 31 +++++++++++++++++++++++++-
 4 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/devito/finite_differences/operators.py b/devito/finite_differences/operators.py
index c19a9cdfce..17452ae664 100644
--- a/devito/finite_differences/operators.py
+++ b/devito/finite_differences/operators.py
@@ -165,13 +165,19 @@ def diag(func, size=None):
         size of the diagonal matrix (size x size).
         Defaults to the number of spatial dimensions when unspecified
     """
+    from devito.types.tensor import TensorFunction, TensorTimeFunction
+    if isinstance(func, TensorFunction):
+        if func.is_TensorValued:
+            return func._new(*func.shape, lambda i, j: func[i, i] if i == j else 0)
+        else:
+            n = func.shape[0]
+            return func._new(n, n, lambda i, j: func[i] if i == j else 0)
+
     dim = size or len(func.dimensions)
     dim = dim-1 if func.is_TimeDependent else dim
     to = getattr(func, 'time_order', 0)
 
-    from devito.types.tensor import TensorFunction, TensorTimeFunction
     tens_func = TensorTimeFunction if func.is_TimeDependent else TensorFunction
-
     comps = [[func if i == j else 0 for i in range(dim)] for j in range(dim)]
     return tens_func(name='diag', grid=func.grid, space_order=func.space_order,
                      components=comps, time_order=to, diagonal=True)
diff --git a/devito/types/basic.py b/devito/types/basic.py
index 593e3ffff7..a28ca2f486 100644
--- a/devito/types/basic.py
+++ b/devito/types/basic.py
@@ -1492,7 +1492,9 @@ def name(self):
     def _rebuild(self, *args, **kwargs):
         # Plain `func` call (row, col, comps)
         if not kwargs.keys() & self.__rkwargs__:
-            assert len(args) == 3
+            if len(args) != 3:
+                raise ValueError("Invalid number of arguments, expected nrow, ncol, "
+                                 "list of components")
             return self._new(*args, **kwargs)
         # We need to rebuild the components with the new name then
         # rebuild the matrix
diff --git a/devito/types/dense.py b/devito/types/dense.py
index d763ea5b71..3eaa6c638d 100644
--- a/devito/types/dense.py
+++ b/devito/types/dense.py
@@ -1089,7 +1089,7 @@ def __staggered_setup__(cls, dimensions, staggered=None, **kwargs):
                     processed.append(sympy.S.NegativeOne)
                 else:
                     processed.append(sympy.S.Zero)
-        return tuple(processed)
+        return Staggering(*processed, getters=dimensions)
 
     @classmethod
     def __indices_setup__(cls, *args, **kwargs):
@@ -1109,9 +1109,8 @@ def __indices_setup__(cls, *args, **kwargs):
             staggered_indices = tuple(args)
         else:
             if not staggered:
-                staggered_indices = (d for d in dimensions)
+                staggered_indices = dimensions
             else:
-                # Staggered indices
                 staggered_indices = (d + i * d.spacing / 2
                                      for d, i in zip(dimensions, staggered))
         return tuple(dimensions), tuple(staggered_indices)
@@ -1119,16 +1118,11 @@ def __indices_setup__(cls, *args, **kwargs):
     @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)
+        return self._staggered
 
     @property
     def is_Staggered(self):
-        if not self.staggered:
-            return False
-        return True
+        return bool(self.staggered)
 
     @classmethod
     def __shape_setup__(cls, **kwargs):
diff --git a/tests/test_tensors.py b/tests/test_tensors.py
index 268bd9ea5c..4cccd73f29 100644
--- a/tests/test_tensors.py
+++ b/tests/test_tensors.py
@@ -5,7 +5,9 @@
 import pytest
 
 from devito import VectorFunction, TensorFunction, VectorTimeFunction, TensorTimeFunction
-from devito import Grid, Function, TimeFunction, Dimension, Eq, div, grad, curl, laplace
+from devito import (
+    Grid, Function, TimeFunction, Dimension, Eq, div, grad, curl, laplace, diag
+)
 from devito.symbolics import retrieve_derivatives
 from devito.types import NODE
 
@@ -465,3 +467,30 @@ def test_rebuild(func1):
         assert j.name == i.name
         assert j.grid == i.grid
         assert j.dimensions == tuple(new_dims)
+
+
+@pytest.mark.parametrize('func1', [Function, TimeFunction,
+                                   TensorFunction, TensorTimeFunction,
+                                   VectorFunction, VectorTimeFunction])
+def test_diag(func1):
+    grid = Grid(tuple([5]*3))
+    f1 = func1(name="f1", grid=grid)
+
+    f2 = diag(f1)
+    assert isinstance(f2, TensorFunction)
+    if f1.is_TimeDependent:
+        assert f2.is_TimeDependent
+    print(f2)
+    assert f2.shape == (3, 3)
+    # Vector input
+    if isinstance(f1, VectorFunction):
+        assert all(f2[i, i] == f1[i] for i in range(3))
+        assert all(f2[i, j] == 0 for i in range(3) for j in range(3) if i != j)
+    # Tensor input
+    elif isinstance(f1, TensorFunction):
+        assert all(f2[i, i] == f1[i, i] for i in range(3))
+        assert all(f2[i, j] == 0 for i in range(3) for j in range(3) if i != j)
+    # Function input
+    else:
+        assert all(f2[i, j] == 0 for i in range(3) for j in range(3) if i != j)
+        assert all(f2[i, i] == f1 for i in range(3))