api: only allow per-component kwargs as Matrix to avoid conflicts

Author: mloubout

devito/types/tensor.py

@@ -80,26 +80,15 @@ def __init_finalize__(self, *args, **kwargs):
         self._space_dimensions = inds
 
     @classmethod
-    def _component_kwargs(cls, *inds, **kwargs):
+    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 k in ('staggered', 'name', 'dimensions', 'shape'):
-                # Standard Function kwargs
-                kw[k] = v
-            elif isinstance(v, MatrixBase):
-                if len(inds) > 1:
-                    kw[k] = v[inds[0], inds[1]]
-                else:
-                    kw[k] = v[inds[0]]
-            elif isinstance(v, (list, tuple)):
-                if len(inds) > 1:
-                    kw[k] = v[inds[0]][inds[1]]
-                else:
-                    kw[k] = v[inds[0]]
+            if isinstance(v, MatrixBase):
+                kw[k] = v[inds]
             else:
                 kw[k] = v
         return kw
@@ -135,7 +124,7 @@ def __subfunc_setup__(cls, *args, **kwargs):
             for j in range(start, stop):
                 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 = 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)

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
 
@@ -498,6 +498,6 @@ def test_diag(func1):
 
 @pytest.mark.parametrize('func1', [TensorFunction, VectorFunction])
 def test_kwargs(func1):
-    orders = [[1, 2], [3, 4]] if func1 is TensorFunction else [1, 2]
+    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 == sympy.Matrix(orders)
+    assert f.space_order == orders