Merge pull request #2583 from devitocodes/tens-rebuild

api: Fix staggering setup and tensor rebuilding

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