api: fix custom coeffs with cross derivatives

Author: mloubout

devito/finite_differences/derivative.py

@@ -484,7 +484,7 @@ def _eval_fd(self, expr, **kwargs):
             assert self.method == 'FD'
             res = cross_derivative(expr, self.dims, self.fd_order, self.deriv_order,
                                    matvec=self.transpose, x0=x0_deriv, expand=expand,
-                                   side=self.side)
+                                   side=self.side, weights=self.weights)
         else:
             assert self.method == 'FD'
             res = generic_derivative(expr, self.dims[0], self.fd_order[0],

devito/finite_differences/tools.py

@@ -92,8 +92,8 @@ def diff_f(expr, deriv_order, dims, fd_order, side=None, **kwargs):
         dims = as_tuple(dims)
         deriv_order = as_tuple(deriv_order)
         fd_order = as_tuple(fd_order)
-        for (d, do, fo) in zip(dims, deriv_order, fd_order):
-            expr = Derivative(expr, d, deriv_order=do, fd_order=fo, side=side, **kwargs)
+        expr = Derivative(expr, *dims, deriv_order=deriv_order, fd_order=fd_order,
+                          side=side, **kwargs)
         return expr
 
     all_combs = dim_with_order(dims, orders)
@@ -331,7 +331,11 @@ def process_weights(weights, expr, dim):
     elif isinstance(weights, Function):
         if len(weights.dimensions) == 1:
             return weights.shape[0], weights.dimensions[0], False
-        wdim = {d for d in weights.dimensions if d not in expr.dimensions}
+        try:
+            # Already a derivative
+            wdim = {d for d in weights.dimensions if d not in expr.base.dimensions}
+        except AttributeError:
+            wdim = {d for d in weights.dimensions if d not in expr.dimensions}
         assert len(wdim) == 1
         wdim = wdim.pop()
         shape = weights.shape

tests/test_symbolic_coefficients.py

@@ -81,7 +81,7 @@ def test_function_coefficients(self):
 
         assert np.all(np.isclose(f0.data[:] - f1.data[:], 0.0, atol=1e-5, rtol=0))
 
-    def test_function_coefficients_xderiv(self):
+    def test_function_coefficients_xderiv_legacy(self):
         p = Dimension('p')
 
         nstc = 8
@@ -105,6 +105,21 @@ def test_function_coefficients_xderiv(self):
         op = Operator(eqn)
         op()
 
+    @pytest.mark.parametrize('order', [2, 4, 6, 8])
+    def test_function_coefficients_xderiv(self, order):
+        p = Dimension('p')
+
+        grid = Grid(shape=(51, 51, 51))
+        x, y, z = grid.dimensions
+
+        f = Function(name='f', grid=grid, space_order=order)
+        w = Function(name='w', space_order=0, shape=(*grid.shape, order + 1),
+                     dimensions=(x, y, z, p))
+
+        expr0 = f.dx(w=w).dy(w=w).evaluate
+        expr1 = f.dxdy(w=w).evaluate
+        assert sp.simplify(expr0 - expr1) == 0
+
     def test_coefficients_w_xreplace(self):
         """Test custom coefficients with an xreplace before they are applied"""
         grid = Grid(shape=(4, 4))