api: allow expression such as f(x)/h_x for custom weights

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)

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
@@ -326,9 +326,13 @@ 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
         try:

tests/test_symbolic_coefficients.py

@@ -120,6 +120,20 @@ def test_function_coefficients_xderiv(self, order):
         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))