api: lif tensor component kwargs init to its own method

Author: mloubout

devito/types/tensor.py

@@ -79,6 +79,25 @@ 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]
+            elif isinstance(v, (list, tuple)):
+                if len(inds) > 1:
+                    kw[k] = v[inds[0]][inds[1]]
+                else:
+                    kw[k] = v[inds[0]]
+            else:
+                kw[k] = v
+        return kw
+
     @classmethod
     def __subfunc_setup__(cls, *args, **kwargs):
         """
@@ -108,20 +127,10 @@ 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):
-                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
+                              'staggered': (stagg[i][j] if stagg is not None
+                                            else (NODE if i == j else (d, dims[j])))}
+                sub_kwargs.update(cls._component_kwargs(i, j, **kwargs))
                 funcs2[j] = cls._sub_type(**sub_kwargs)
             funcs.append(funcs2)
 
@@ -335,17 +344,9 @@ 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
-            # 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
+            sub_kwargs = {'name': f"{name}_{d.name}",
+                          'staggered': stagg[i] if stagg is not None else d}
+            sub_kwargs.update(cls._component_kwargs(i, **kwargs))
             funcs.append(cls._sub_type(**sub_kwargs))
 
         return funcs