Merge pull request #2672 from devitocodes/xderiv-fix

api: fix custom coeffs with cross derivatives

Author: mloubout

devito/finite_differences/derivative.py

@@ -213,7 +213,7 @@ def _process_weights(cls, **kwargs):
         weights = kwargs.get('weights', kwargs.get('w'))
         if weights is None:
             return None
-        elif isinstance(weights, sympy.Function):
+        elif isinstance(weights, Differentiable):
             return weights
         else:
             return as_tuple(weights)
@@ -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/finite_difference.py

@@ -59,7 +59,7 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, side=None, **kw
     Semantically, this is equivalent to
 
     >>> (f*g).dxdy
-    Derivative(Derivative(f(x, y)*g(x, y), x), y)
+    Derivative(f(x, y)*g(x, y), x, y)
 
     The only difference is that in the latter case derivatives remain unevaluated.
     The expanded form is obtained via ``evaluate``
@@ -158,7 +158,9 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
     # `coefficients` method (`taylor` or `symbolic`)
     if weights is None:
         weights = fd_weights_registry[coefficients](expr, deriv_order, indices, x0)
-    elif wdim is not None:
+    # Did fd_weights_registry return a new Function/Expression instead of a values?
+    _, wdim, _ = process_weights(weights, expr, dim)
+    if wdim is not None:
         weights = [weights._subs(wdim, i) for i in range(len(indices))]
 
     # Enforce fixed precision FD coefficients to avoid variations in results

devito/finite_differences/tools.py

@@ -2,7 +2,7 @@
 from itertools import product
 
 import numpy as np
-from sympy import S, finite_diff_weights, cacheit, sympify, Function, Rational
+from sympy import S, finite_diff_weights, cacheit, sympify, Rational, Expr
 
 from devito.logger import warning
 from devito.tools import Tag, as_tuple
@@ -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)
@@ -326,12 +326,20 @@ def make_shift_x0(shift, ndim):
 
 
 def process_weights(weights, expr, dim):
+    from devito.symbolics import retrieve_functions
     if weights is None:
         return 0, None, False
-    elif isinstance(weights, Function):
+    elif isinstance(weights, Expr):
+        w_func = retrieve_functions(weights)
+        assert len(w_func) == 1, "Only one function expected in weights"
+        weights = w_func[0]
         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_derivatives.py

@@ -87,19 +87,19 @@ def test_stencil_derivative(self, SymbolType, dim):
 
     @pytest.mark.parametrize('SymbolType, derivative, dim, expected', [
         (Function, ['dx2'], 3, 'Derivative(u(x, y, z), (x, 2))'),
-        (Function, ['dx2dy'], 3, 'Derivative(Derivative(u(x, y, z), (x, 2)), y)'),
+        (Function, ['dx2dy'], 3, 'Derivative(u(x, y, z), (x, 2), y)'),
         (Function, ['dx2dydz'], 3,
-         'Derivative(Derivative(Derivative(u(x, y, z), (x, 2)), y), z)'),
+         'Derivative(u(x, y, z), (x, 2), y, z)'),
         (Function, ['dx2', 'dy'], 3, 'Derivative(Derivative(u(x, y, z), (x, 2)), y)'),
         (Function, ['dx2dy', 'dz2'], 3,
-         'Derivative(Derivative(Derivative(u(x, y, z), (x, 2)), y), (z, 2))'),
+         'Derivative(Derivative(u(x, y, z), (x, 2), y), (z, 2))'),
         (TimeFunction, ['dx2'], 3, 'Derivative(u(t, x, y, z), (x, 2))'),
-        (TimeFunction, ['dx2dy'], 3, 'Derivative(Derivative(u(t, x, y, z), (x, 2)), y)'),
+        (TimeFunction, ['dx2dy'], 3, 'Derivative(u(t, x, y, z), (x, 2), y)'),
         (TimeFunction, ['dx2', 'dy'], 3,
          'Derivative(Derivative(u(t, x, y, z), (x, 2)), y)'),
         (TimeFunction, ['dx', 'dy', 'dx2', 'dz', 'dydz'], 3,
-         'Derivative(Derivative(Derivative(Derivative(Derivative(Derivative(' +
-         'u(t, x, y, z), x), y), (x, 2)), z), y), z)')
+         'Derivative(Derivative(Derivative(Derivative(Derivative(' +
+         'u(t, x, y, z), x), y), (x, 2)), z), y, z)')
     ])
     def test_unevaluation(self, SymbolType, derivative, dim, expected):
         u = SymbolType(name='u', grid=self.grid, time_order=2, space_order=2)
@@ -112,13 +112,13 @@ def test_unevaluation(self, SymbolType, derivative, dim, expected):
 
     @pytest.mark.parametrize('expr,expected', [
         ('u.dx + u.dy', 'Derivative(u, x) + Derivative(u, y)'),
-        ('u.dxdy', 'Derivative(Derivative(u, x), y)'),
+        ('u.dxdy', 'Derivative(u, x, y)'),
         ('u.laplace',
          'Derivative(u, (x, 2)) + Derivative(u, (y, 2)) + Derivative(u, (z, 2))'),
         ('(u.dx + u.dy).dx', 'Derivative(Derivative(u, x) + Derivative(u, y), x)'),
         ('((u.dx + u.dy).dx + u.dxdy).dx',
          'Derivative(Derivative(Derivative(u, x) + Derivative(u, y), x) +' +
-         ' Derivative(Derivative(u, x), y), x)'),
+         ' Derivative(u, x, y), x)'),
         ('(u**4).dx', 'Derivative(u**4, x)'),
         ('(u/4).dx', 'Derivative(u/4, x)'),
         ('((u.dx + v.dy).dx * v.dx).dy.dz',
@@ -741,7 +741,7 @@ def test_cross_newnest(self):
         grid = Grid((11, 11))
         f = Function(name="f", grid=grid, space_order=2)
 
-        assert f.dxdy == f.dx.dy
+        assert simplify((f.dxdy - f.dx.dy).evaluate) == 0
 
     def test_nested_call(self):
         grid = Grid((11, 11))

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,35 @@ 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_expr(self):
+        p = Dimension('p')
+
+        grid = Grid(shape=(51, 51, 51))
+        x, y, z = grid.dimensions
+
+        f = Function(name='f', grid=grid, space_order=4)
+        w = Function(name='w', space_order=0, shape=(*grid.shape, 5),
+                     dimensions=(x, y, z, p))
+
+        expr0 = f.dx(w=w/x.spacing).evaluate
+        expr1 = f.dx(w=w).evaluate / x.spacing
+        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))