api: fix derivatives kw for tensors

Author: mloubout

devito/operations/interpolators.py

@@ -125,6 +125,8 @@ def __repr__(self):
         return (f"Interpolation({repr(self.expr)} into "
                 f"{repr(self.interpolator.sfunction)})")
 
+    __str__ = __repr__
+
 
 class Injection(UnevaluatedSparseOperation):
 
@@ -152,6 +154,8 @@ def operation(self, **kwargs):
     def __repr__(self):
         return f"Injection({repr(self.expr)} into {repr(self.field)})"
 
+    __str__ = __repr__
+
 
 class GenericInterpolator(ABC):
 

devito/types/basic.py

@@ -1574,6 +1574,16 @@ def adjoint(self, inner=True):
         # Real valued adjoint is transpose
         return self.transpose(inner=inner)
 
+    def __call__(self, **kwargs):
+        """
+        Derivative custom inputs (weights/x0/...) is done through call
+        and needs to be applied to each component through applyfunc
+        """
+        try:
+            return self.applyfunc(lambda x: x(**kwargs))
+        except TypeError as e:
+            raise f"{self.name} not callable with {kwargs}" from e
+
     @call_highest_priority('__radd__')
     def __add__(self, other):
         try:

tests/test_tensors.py

@@ -442,6 +442,21 @@ def test_custom_coeffs_tensor():
             assert list(drv.weights) == c
 
 
+@pytest.mark.parametrize('func', [TensorFunction, TensorTimeFunction,
+                                  VectorFunction, VectorTimeFunction])
+def test_custom_coeffs_tensor_basic(func):
+    grid = Grid(tuple([5]*3))
+    f = func(name="t", grid=grid, space_order=2)
+
+    # Custom coefficients
+    c = [10, 10, 10]
+
+    df = f.dx(w=c)
+    for (fi, dfi) in zip(f.values(), df.values()):
+        assert dfi == fi.dx(w=c)
+        assert list(dfi.weights) == c
+
+
 @pytest.mark.parametrize('func1', [TensorFunction, TensorTimeFunction,
                                    VectorFunction, VectorTimeFunction])
 def test_rebuild(func1):