api: fix kwargs processing for tensor functions

Author: mloubout

devito/builtins/initializers.py

@@ -79,6 +79,7 @@ def assign(f, rhs=0, options=None, name='assign', assign_halo=False, **kwargs):
         eqs = [eq.xreplace(subs) for eq in eqs]
 
     op = dv.Operator(eqs, name=name, **kwargs)
+
     try:
         op()
     except ValueError:

devito/types/basic.py

@@ -1489,6 +1489,15 @@ def __subfunc_setup__(cls, *args, **kwargs):
         """Setup each component of the tensor as a Devito type."""
         return []
 
+    @classmethod
+    def _sympify(self, arg):
+        # This is used internally by sympy to process arguments at rebuilt. And since
+        # some of our properties are non-sympyfiable we need to have a fallback
+        try:
+            return super()._sympify(arg)
+        except sympy.SympifyError:
+            return arg
+
     @property
     def grid(self):
         """

devito/types/dense.py

@@ -1019,7 +1019,7 @@ def __init_finalize__(self, *args, **kwargs):
 
         # Space order
         space_order = kwargs.get('space_order', 1)
-        if isinstance(space_order, int):
+        if is_integer(space_order):
             self._space_order = space_order
         elif isinstance(space_order, tuple) and len(space_order) >= 2:
             self._space_order = space_order[0]
@@ -1175,7 +1175,7 @@ def __halo_setup__(self, **kwargs):
                 halo = tuple(halo[d] for d in self.dimensions)
         else:
             space_order = kwargs.get('space_order', 1)
-            if isinstance(space_order, int):
+            if is_integer(space_order):
                 v = (space_order, space_order)
                 halo = [v if i.is_Space else (0, 0) for i in self.dimensions]
 
@@ -1208,12 +1208,12 @@ def __padding_setup__(self, **kwargs):
         elif isinstance(padding, DimensionTuple):
             padding = tuple(padding[d] for d in self.dimensions)
 
-        elif isinstance(padding, int):
+        elif is_integer(padding):
             padding = tuple((0, padding) if d.is_Space else (0, 0)
                             for d in self.dimensions)
 
         elif isinstance(padding, tuple) and len(padding) == self.ndim:
-            padding = tuple((0, i) if isinstance(i, int) else i for i in padding)
+            padding = tuple((0, i) if is_integer(i) else i for i in padding)
 
         else:
             raise TypeError("`padding` must be int or %d-tuple of ints" % self.ndim)
@@ -1398,7 +1398,7 @@ def __init_finalize__(self, *args, **kwargs):
         self._time_order = kwargs.get('time_order', 1)
         super().__init_finalize__(*args, **kwargs)
 
-        if not isinstance(self.time_order, int):
+        if not is_integer(self.time_order):
             raise TypeError("`time_order` must be int")
 
         self.save = kwargs.get('save')
@@ -1420,7 +1420,7 @@ def __indices_setup__(cls, *args, **kwargs):
             time_dim = kwargs.get('time_dim')
 
             if time_dim is None:
-                time_dim = grid.time_dim if isinstance(save, int) else grid.stepping_dim
+                time_dim = grid.time_dim if is_integer(save) else grid.stepping_dim
             elif not (isinstance(time_dim, Dimension) and time_dim.is_Time):
                 raise TypeError("`time_dim` must be a time dimension")
             dimensions = list(Function.__indices_setup__(**kwargs)[0])
@@ -1450,7 +1450,7 @@ def __shape_setup__(cls, **kwargs):
                 shape.insert(cls._time_position, time_order + 1)
             elif isinstance(save, Buffer):
                 shape.insert(cls._time_position, save.val)
-            elif isinstance(save, int):
+            elif is_integer(save):
                 shape.insert(cls._time_position, save)
             else:
                 raise TypeError("`save` can be None, int or Buffer, not %s" % type(save))

devito/types/tensor.py

@@ -2,6 +2,7 @@
 from functools import cached_property
 
 import numpy as np
+from sympy.matrices.matrixbase import MatrixBase
 from sympy.core.sympify import converter as sympify_converter
 
 from devito.finite_differences import Differentiable
@@ -107,10 +108,21 @@ 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):
-                kwargs["name"] = "%s_%s%s" % (name, d.name, dims[j].name)
-                kwargs["staggered"] = (stagg[i][j] if stagg is not None
-                                       else (NODE if i == j else (d, dims[j])))
-                funcs2[j] = cls._sub_type(**kwargs)
+                staggj = (stagg[i][j] if stagg is not None
+                          else (NODE if i == j else (d, dims[j])))
+                # Setup kwargs for subfunction
+                # Through rebuilding or user input, the kwargs could be
+                # Tensors as well from a per-component property
+                sub_kwargs = {'name': f"{name}_{d.name}{dims[j].name}",
+                              'staggered': staggj}
+                for k, v in kwargs.items():
+                    if isinstance(v, MatrixBase):
+                        sub_kwargs[k] = v[i, j]
+                    elif isinstance(v, (list, tuple)):
+                        sub_kwargs[k] = v[i][j]
+                    else:
+                        sub_kwargs[k] = v
+                funcs2[j] = cls._sub_type(**sub_kwargs)
             funcs.append(funcs2)
 
         # Symmetrize and fill diagonal if symmetric
@@ -169,7 +181,11 @@ def root_dimensions(self):
     @cached_property
     def space_order(self):
         """The space order for all components."""
-        return ({a.space_order for a in self} - {None}).pop()
+        orders = self.applyfunc(lambda x: x.space_order)
+        if len(set(orders)) > 1:
+            return orders
+        else:
+            return orders[0]
 
     @property
     def is_diagonal(self):
@@ -321,7 +337,16 @@ def __subfunc_setup__(cls, *args, **kwargs):
         for i, d in enumerate(dims):
             kwargs["name"] = "%s_%s" % (name, d.name)
             kwargs["staggered"] = stagg[i] if stagg is not None else d
-            funcs.append(cls._sub_type(**kwargs))
+            # Setup kwargs for subfunction
+            # Through rebuilding or user input, the kwargs could be
+            # Tensors as well from a per-component property
+            sub_kwargs = {}
+            for k, v in kwargs.items():
+                if isinstance(v, (list, tuple, MatrixBase)):
+                    sub_kwargs[k] = v[i]
+                else:
+                    sub_kwargs[k] = v
+            funcs.append(cls._sub_type(**sub_kwargs))
 
         return funcs
 

tests/test_tensors.py

@@ -494,3 +494,10 @@ def test_diag(func1):
     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))
+
+
+@pytest.mark.parametrize('func1', [TensorFunction, VectorFunction])
+def test_kwargs(func1):
+    orders = [[1, 2], [3, 4]] if func1 is TensorFunction else [1, 2]
+    f = func1(name="f", grid=Grid((5, 5)), space_order=orders, symmetric=False)
+    assert f.space_order == sympy.Matrix(orders)