Merge pull request #2604 from devitocodes/tens-kwargs

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
@@ -78,6 +79,20 @@ def __init_finalize__(self, *args, **kwargs):
         inds, _ = Function.__indices_setup__(grid=grid, dimensions=dimensions)
         self._space_dimensions = inds
 
+    @classmethod
+    def _component_kwargs(cls, inds, **kwargs):
+        """
+        Get the kwargs for a single component
+        from the kwargs of the TensorFunction.
+        """
+        kw = {}
+        for k, v in kwargs.items():
+            if isinstance(v, MatrixBase):
+                kw[k] = v[inds]
+            else:
+                kw[k] = v
+        return kw
+
     @classmethod
     def __subfunc_setup__(cls, *args, **kwargs):
         """
@@ -107,10 +122,12 @@ 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])))
+                sub_kwargs = cls._component_kwargs((i, j), **kwargs)
+                sub_kwargs.update({'name': f"{name}_{d.name}{dims[j].name}",
+                                   'staggered': staggj})
+                funcs2[j] = cls._sub_type(**sub_kwargs)
             funcs.append(funcs2)
 
         # Symmetrize and fill diagonal if symmetric
@@ -169,7 +186,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):
@@ -319,9 +340,10 @@ def __subfunc_setup__(cls, *args, **kwargs):
         stagg = kwargs.get("staggered", None)
         name = kwargs.get("name")
         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))
+            sub_kwargs = cls._component_kwargs(i, **kwargs)
+            sub_kwargs.update({'name': f"{name}_{d.name}",
+                               'staggered': stagg[i] if stagg is not None else d})
+            funcs.append(cls._sub_type(**sub_kwargs))
 
         return funcs
 

docker/Dockerfile.devito

@@ -13,11 +13,11 @@ ARG GROUP_ID=1000
 
 ##################  Install devito ############################################
 
+ENV PIP_USE_PEP517=1
 # Install pip dependencies
 RUN python3 -m venv /venv && \
-    /venv/bin/pip install --no-cache-dir --upgrade pip && \
+    /venv/bin/pip install --no-cache-dir --upgrade pip wheel setuptools && \
     /venv/bin/pip install --no-cache-dir jupyter && \
-    /venv/bin/pip install --no-cache-dir --upgrade wheel setuptools && \
     ln -fs /app/nvtop/build/src/nvtop /venv/bin/nvtop
 
 # Copy Devito

tests/test_tensors.py

@@ -1,6 +1,6 @@
 import numpy as np
 import sympy
-from sympy import Rational
+from sympy import Rational, Matrix
 
 import pytest
 
@@ -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 = Matrix([[1, 2], [3, 4]]) if func1 is TensorFunction else Matrix([1, 2])
+    f = func1(name="f", grid=Grid((5, 5)), space_order=orders, symmetric=False)
+    assert f.space_order == orders