api: add support for diag(tensor) and diag(vector)

Author: mloubout

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)

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

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,26 +1109,20 @@ 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)
 
     @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):

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))