dsl: better fix for custom coefficients on staggered grids

devito/finite_differences/coefficients.py

@@ -6,7 +6,7 @@
 from devito.tools import filter_ordered, as_tuple
 from devito.symbolics.search import retrieve_dimensions
 
-__all__ = ['Coefficient', 'Substitutions', 'default_rules']
+__all__ = ['Coefficient', 'Substitutions', 'default_rules', 'all_rules']
 
 
 class Coefficient(object):
@@ -227,6 +227,34 @@ def generate_subs(i):
 
         return rules
 
+    def update_rules(self, obj):
+        """Update the specified rules to reflect staggering in an equation"""
+        # Determine which 'rules' are expected
+        sym = get_sym(self._function_list)
+        terms = obj.find(sym)
+        args_expected = filter_ordered(term.args[1:] for term in terms)
+        args_expected_dim = [(arg[0], arg[1], retrieve_dimensions(arg[2])[0])
+                             for arg in args_expected]
+
+        # Modify dictionary keys where expected index does not match index in rules
+        rules = self._rules.copy()  # Get a copy to modify, to preserve base rules
+        for rule in self._rules:
+            rule_arg = rule.args[1:]
+            rule_arg_dim = (rule_arg[0], rule_arg[1],
+                            retrieve_dimensions(rule_arg[2])[0])
+            if rule_arg_dim in args_expected_dim and rule_arg not in args_expected:
+                # Rule matches expected in terms of dimensions, but index is
+                # mismatched (due to staggering of equation)
+
+                # Find index in args_expected_dim
+                pos = args_expected_dim.index(rule_arg_dim)
+                # Replace key in rules with one using modified index taken from
+                # the expected
+                replacement = rule.args[:-1] + (args_expected[pos][-1],)
+                rules[sym(*replacement)] = rules.pop(rule)
+
+        return rules
+
 
 def default_rules(obj, functions):
 
@@ -246,7 +274,8 @@ def generate_subs(deriv_order, function, index):
         mapper = {dim: index}
 
         indices, x0 = generate_indices(function, dim,
-                                       fd_order, side=None, x0=mapper)
+                                       fd_order, side=None, x0=mapper,
+                                       symbolic=True)
 
         coeffs = sympy.finite_diff_weights(deriv_order, indices, x0)[-1][-1]
 
@@ -262,9 +291,10 @@ def generate_subs(deriv_order, function, index):
     args_present = filter_ordered(term.args[1:] for term in terms)
 
     subs = obj.substitutions
+
     if subs:
-        args_provided = [(i.deriv_order, i.function, i.index)
-                         for i in subs.coefficients]
+        # Check against the updated rules when determining rules the user has provided
+        args_provided = filter_ordered(rule.args[1:] for rule in subs.update_rules(obj))
     else:
         args_provided = []
 
@@ -289,3 +319,17 @@ def get_sym(functions):
             pass
     # Shouldn't arrive here
     raise TypeError("Failed to retreive symbol")
+
+
+def all_rules(obj, functions):
+    """Return all substitution rules for an Eq"""
+    # Default rules
+    d_rules = default_rules(obj, functions)
+    # User rules
+    subs = obj.substitutions
+    if subs:
+        u_rules = subs.update_rules(obj)
+    else:
+        u_rules = {}
+
+    return {**d_rules, **u_rules}

devito/finite_differences/finite_difference.py

@@ -249,13 +249,15 @@ def generic_derivative(expr, dim, fd_order, deriv_order, symbolic=False,
     # first order fd that is a lot better
     if deriv_order == 1 and fd_order == 2 and not symbolic:
         fd_order = 1
-    # Stencil positions
-    indices, x0 = generate_indices(expr, dim, fd_order, x0=x0)
 
     # Finite difference weights from Taylor approximation with these positions
     if symbolic:
+        # Stencil positions
+        indices, x0 = generate_indices(expr, dim, fd_order, x0=x0, symbolic=True)
         c = symbolic_weights(expr, deriv_order, indices, x0)
     else:
+        # Stencil positions
+        indices, x0 = generate_indices(expr, dim, fd_order, x0=x0)
         c = numeric_weights(deriv_order, indices, x0)
 
     return indices_weights_to_fd(expr, dim, indices, c, matvec=matvec.val)

devito/finite_differences/tools.py

@@ -140,7 +140,7 @@ def diff_f(expr, deriv_order, dims, fd_order, side=None, **kwargs):
 
 
 def symbolic_weights(function, deriv_order, indices, dim):
-    return [function._coeff_symbol(indices[j], deriv_order, function, dim)
+    return [function._coeff_symbol(indices[j], deriv_order, function.function, dim)
             for j in range(0, len(indices))]
 
 
@@ -149,7 +149,7 @@ def numeric_weights(deriv_order, indices, x0):
     return finite_diff_weights(deriv_order, indices, x0)[-1][-1]
 
 
-def generate_indices(func, dim, order, side=None, x0=None):
+def generate_indices(func, dim, order, side=None, x0=None, symbolic=False):
     """
     Indices for the finite-difference scheme
 
@@ -165,14 +165,17 @@ def generate_indices(func, dim, order, side=None, x0=None):
         Side of the scheme, (centered, left, right)
     x0: Dict of {Dimension: Dimension or Expr or Number}
         Origin of the scheme, ie. `x`, `x + .5 * x.spacing`, ...
+    symbolic: bool
+        Override standard indices generation for symbolic coefficients
 
     Returns
     -------
     Ordered list of indices
     """
     # If staggered finited difference
     if func.is_Staggered and not dim.is_Time:
-        x0, ind = generate_indices_staggered(func, dim, order, side=side, x0=x0)
+        x0, ind = generate_indices_staggered(func, dim, order, side=side, x0=x0,
+                                             symbolic=symbolic)
     else:
         x0 = (x0 or {dim: dim}).get(dim, dim)
         # Check if called from first_derivative()
@@ -220,7 +223,7 @@ def generate_indices_cartesian(dim, order, side, x0):
     return tuple(ind)
 
 
-def generate_indices_staggered(func, dim, order, side=None, x0=None):
+def generate_indices_staggered(func, dim, order, side=None, x0=None, symbolic=False):
     """
     Indices for the finite-difference scheme on a staggered grid
 
@@ -236,6 +239,8 @@ def generate_indices_staggered(func, dim, order, side=None, x0=None):
         Side of the scheme, (centered, left, right)
     x0: Dict of {Dimension: Dimension or Expr or Number}
         Origin of the scheme, ie. `x`, `x + .5 * x.spacing`, ...
+    symbolic: bool
+        Override standard indices generation for symbolic coefficients
 
     Returns
     -------
@@ -247,6 +252,9 @@ def generate_indices_staggered(func, dim, order, side=None, x0=None):
         ind0 = func.indices_ref[dim]
     except AttributeError:
         ind0 = start
+    if symbolic and order >= 2:
+        ind = [ind0 + i*diff for i in range(-order//2, order//2+1)]
+        return start, tuple(ind)
     if start != ind0:
         ind = [start - diff/2 - i * diff for i in range(0, order//2)][::-1]
         ind += [start + diff/2 + i * diff for i in range(0, order//2)]

devito/types/equation.py

@@ -5,7 +5,7 @@
 
 from cached_property import cached_property
 
-from devito.finite_differences import default_rules
+from devito.finite_differences import all_rules
 from devito.logger import warning
 from devito.tools import as_tuple
 from devito.types.lazy import Evaluable
@@ -95,9 +95,9 @@ def evaluate(self):
         if eq._uses_symbolic_coefficients:
             # NOTE: As Coefficients.py is expanded we will not want
             # all rules to be expunged during this procress.
-            rules = default_rules(eq, eq._symbolic_functions)
+            rules = all_rules(eq, eq._symbolic_functions)
             try:
-                eq = eq.xreplace({**eq.substitutions.rules, **rules})
+                eq = eq.xreplace(rules)
             except AttributeError:
                 if bool(rules):
                     eq = eq.xreplace(rules)

tests/test_symbolic_coefficients.py

@@ -243,6 +243,113 @@ def test_default_rules_deriv_offset(self, so, offset):
         eval_sym = str(Eq(g, f_sym.dx(x0=x+offset*h_x/2)).evaluate.evalf(_PRECISION))
         assert eval_std == eval_sym
 
+    @pytest.mark.parametrize('so, dir, expected', [(2, 'fwd', '1.0*f(x)'
+                                                              ' + 1.0*f(x - h_x)'
+                                                              ' + 1.0*f(x + h_x)'),
+                                                   (4, 'fwd', '1.0*f(x)'
+                                                              ' + 1.0*f(x - 2*h_x)'
+                                                              ' + 1.0*f(x - h_x)'
+                                                              ' + 1.0*f(x + h_x)'
+                                                              ' + 1.0*f(x + 2*h_x)'),
+                                                   (6, 'fwd', '1.0*f(x)'
+                                                              ' + 1.0*f(x - 3*h_x)'
+                                                              ' + 1.0*f(x - 2*h_x)'
+                                                              ' + 1.0*f(x - h_x)'
+                                                              ' + 1.0*f(x + h_x)'
+                                                              ' + 1.0*f(x + 2*h_x)'
+                                                              ' + 1.0*f(x + 3*h_x)'),
+                                                   (2, 'rev', '1.0*f(x - h_x/2)'
+                                                              ' + 1.0*f(x + h_x/2)'
+                                                              ' + 1.0*f(x + 3*h_x/2)'),
+                                                   (4, 'rev', '1.0*f(x - 3*h_x/2)'
+                                                              ' + 1.0*f(x - h_x/2)'
+                                                              ' + 1.0*f(x + h_x/2)'
+                                                              ' + 1.0*f(x + 3*h_x/2)'
+                                                              ' + 1.0*f(x + 5*h_x/2)')])
+    def test_custom_coefficients_staggered_eq(self, so, dir, expected):
+        """Test custom coefficients for staggered equations"""
+        grid = Grid(shape=(11,), extent=(10.,))
+        x = grid.dimensions[0]
+        if dir == 'fwd':
+            f = Function(name='f', grid=grid, space_order=so, staggered=NODE,
+                         coefficients='symbolic')
+            g = Function(name='g', grid=grid, space_order=so, staggered=x,
+                         coefficients='symbolic')
+        elif dir == 'rev':
+            f = Function(name='f', grid=grid, space_order=so, staggered=x,
+                         coefficients='symbolic')
+            g = Function(name='g', grid=grid, space_order=so, staggered=NODE,
+                         coefficients='symbolic')
+        else:
+            raise ValueError("Invalid stagger direction '{}'".format(dir))
+        weights = np.ones(so+1)
+        coeffs = Coefficient(1, f, x, weights)
+        eq = Eq(g, f.dx, coefficients=Substitutions(coeffs))
+        assert str(eq.evaluate.rhs) == expected
+
+    @pytest.mark.parametrize('so, expected',
+                             [(2, '1.0*f(x) + 1.0*f(x - h_x) + 1.0*f(x + h_x)'
+                                  ' - g(x)/h_x + g(x + h_x)/h_x'),
+                              (4, '1.0*f(x) + 1.0*f(x - 2*h_x) + 1.0*f(x - h_x)'
+                                  ' + 1.0*f(x + h_x) + 1.0*f(x + 2*h_x) - 9*g(x)/(8*h_x)'
+                                  ' + g(x - h_x)/(24*h_x) + 9*g(x + h_x)/(8*h_x)'
+                                  ' - g(x + 2*h_x)/(24*h_x)'),
+                              (6, '1.0*f(x) + 1.0*f(x - 3*h_x) + 1.0*f(x - 2*h_x)'
+                                  ' + 1.0*f(x - h_x) + 1.0*f(x + h_x) + 1.0*f(x + 2*h_x)'
+                                  ' + 1.0*f(x + 3*h_x) - 75*g(x)/(64*h_x)'
+                                  ' - 3*g(x - 2*h_x)/(640*h_x)'
+                                  ' + 25*g(x - h_x)/(384*h_x)'
+                                  ' + 75*g(x + h_x)/(64*h_x)'
+                                  ' - 25*g(x + 2*h_x)/(384*h_x)'
+                                  ' + 3*g(x + 3*h_x)/(640*h_x)')])
+    def test_custom_default_combo(self, so, expected):
+        """
+        Test that staggered equations containing a mixture of default and custom
+        coefficients evaluate as expected.
+        """
+        grid = Grid(shape=(11,), extent=(10.,))
+        x = grid.dimensions[0]
+        f = Function(name='f', grid=grid, space_order=so, staggered=NODE,
+                     coefficients='symbolic')
+        g = Function(name='g', grid=grid, space_order=so, staggered=NODE,
+                     coefficients='symbolic')
+        h = Function(name='h', grid=grid, space_order=so, staggered=x,
+                     coefficients='symbolic')
+
+        weights = np.ones(so+1)
+        coeffs = Coefficient(1, f, x, weights)
+        eq = Eq(h, f.dx + g.dx, coefficients=Substitutions(coeffs))
+        assert str(eq.evaluate.rhs) == expected
+
+    @pytest.mark.parametrize('so', [2, 4, 6])
+    @pytest.mark.parametrize('offset', [1, -1])
+    def test_custom_deriv_offset(self, so, offset):
+        """
+        Test that stencils with custom coeffs evaluate correctly when derivatives
+        are evaluated at offset x0.
+        """
+        grid = Grid(shape=(11,), extent=(10.,))
+        x = grid.dimensions[0]
+        h_x = x.spacing
+        f = Function(name='f', grid=grid, space_order=so, coefficients='symbolic')
+        g = Function(name='g', grid=grid, space_order=so)
+
+        weights = np.arange(so+1) + 1
+        coeffs = Coefficient(1, f, x, weights)
+
+        bench_weights = np.arange(so+1) + 1
+        if offset == 1:
+            bench_weights[0] = 0
+        else:
+            bench_weights[-1] = 0
+        bench_coeffs = Coefficient(1, f, x, bench_weights)
+
+        eq = Eq(g, f.dx(x0=x+offset*h_x/2),
+                coefficients=Substitutions(coeffs))
+        eq2 = Eq(g, f.dx, coefficients=Substitutions(bench_coeffs))
+
+        assert str(eq.evaluate.evalf(_PRECISION)) == str(eq2.evaluate.evalf(_PRECISION))
+
     @pytest.mark.parametrize('order', [2, 4, 6])
     def test_staggered_array(self, order):
         """Test custom coefficients provided as an array on a staggered grid"""
@@ -345,3 +452,91 @@ def test_with_timefunction(self, stagger):
         Operator([eq_f, eq_g])(t_m=0, t_M=1)
 
         assert np.allclose(f.data[-1], -g.data[-1], atol=1e-7)
+
+    @pytest.mark.parametrize('order, expected', [(2, '1.0*f(t, x) + 1.0*f(t, x - h_x)'
+                                                     ' + 1.0*f(t, x + h_x)'
+                                                     ' - f(t, x - h_x)/(2*h_x)'
+                                                     ' + f(t, x + h_x)/(2*h_x)'),
+                                                 (4, '1.0*f(t, x) + 1.0*f(t, x - 2*h_x)'
+                                                     ' + 1.0*f(t, x - h_x)'
+                                                     ' + 1.0*f(t, x + h_x)'
+                                                     ' + 1.0*f(t, x + 2*h_x)'
+                                                     ' + f(t, x - 2*h_x)/(12*h_x)'
+                                                     ' - 2*f(t, x - h_x)/(3*h_x)'
+                                                     ' + 2*f(t, x + h_x)/(3*h_x)'
+                                                     ' - f(t, x + 2*h_x)/(12*h_x)'),
+                                                 (6, '1.0*f(t, x) + 1.0*f(t, x - 3*h_x)'
+                                                     ' + 1.0*f(t, x - 2*h_x)'
+                                                     ' + 1.0*f(t, x - h_x)'
+                                                     ' + 1.0*f(t, x + h_x)'
+                                                     ' + 1.0*f(t, x + 2*h_x)'
+                                                     ' + 1.0*f(t, x + 3*h_x)'
+                                                     ' - f(t, x - 3*h_x)/(60*h_x)'
+                                                     ' + 3*f(t, x - 2*h_x)/(20*h_x)'
+                                                     ' - 3*f(t, x - h_x)/(4*h_x)'
+                                                     ' + 3*f(t, x + h_x)/(4*h_x)'
+                                                     ' - 3*f(t, x + 2*h_x)/(20*h_x)'
+                                                     ' + f(t, x + 3*h_x)/(60*h_x)')])
+    def test_some_substitutions(self, order, expected):
+        """
+        Test that equations where only some derivatives have substitutions evaluate
+        correctly.
+        """
+        grid = Grid(shape=(11,), extent=(10.,))
+        x = grid.dimensions[0]
+
+        f = TimeFunction(name='f', grid=grid, space_order=order, coefficients='symbolic')
+
+        weights = np.ones(order+1)
+        coeffs = Coefficient(2, f, x, weights)
+
+        eq = Eq(f.dx + f.dx2, 0, coefficients=Substitutions(coeffs))
+
+        assert str(eq.evaluate.lhs) == expected
+
+    @pytest.mark.parametrize('order, expected', [(2, '1.0*f(t - dt, x)'
+                                                     ' + 1.0*f(t - dt, x - h_x)'
+                                                     ' + 1.0*f(t - dt, x + h_x)'
+                                                     ' + 1.0*f(t + dt, x)'
+                                                     ' + 1.0*f(t + dt, x - h_x)'
+                                                     ' + 1.0*f(t + dt, x + h_x)'),
+                                                 (4, '1.0*f(t - dt, x)'
+                                                     ' + 1.0*f(t - dt, x - 2*h_x)'
+                                                     ' + 1.0*f(t - dt, x - h_x)'
+                                                     ' + 1.0*f(t - dt, x + h_x)'
+                                                     ' + 1.0*f(t - dt, x + 2*h_x)'
+                                                     ' + 1.0*f(t + dt, x)'
+                                                     ' + 1.0*f(t + dt, x - 2*h_x)'
+                                                     ' + 1.0*f(t + dt, x - h_x)'
+                                                     ' + 1.0*f(t + dt, x + h_x)'
+                                                     ' + 1.0*f(t + dt, x + 2*h_x)'),
+                                                 (6, '1.0*f(t - dt, x)'
+                                                     ' + 1.0*f(t - dt, x - 3*h_x)'
+                                                     ' + 1.0*f(t - dt, x - 2*h_x)'
+                                                     ' + 1.0*f(t - dt, x - h_x)'
+                                                     ' + 1.0*f(t - dt, x + h_x)'
+                                                     ' + 1.0*f(t - dt, x + 2*h_x)'
+                                                     ' + 1.0*f(t - dt, x + 3*h_x)'
+                                                     ' + 1.0*f(t + dt, x)'
+                                                     ' + 1.0*f(t + dt, x - 3*h_x)'
+                                                     ' + 1.0*f(t + dt, x - 2*h_x)'
+                                                     ' + 1.0*f(t + dt, x - h_x)'
+                                                     ' + 1.0*f(t + dt, x + h_x)'
+                                                     ' + 1.0*f(t + dt, x + 2*h_x)'
+                                                     ' + 1.0*f(t + dt, x + 3*h_x)')])
+    def test_forward_backward_timestep(self, order, expected):
+        """
+        Test to ensure that substitutions are applied to derivatives taken at both
+        forward and backward timesteps.
+        """
+        grid = Grid(shape=(11,), extent=(10.,))
+        x = grid.dimensions[0]
+
+        f = TimeFunction(name='f', grid=grid, space_order=order, coefficients='symbolic')
+
+        weights = np.ones(order+1)
+        coeffs = Coefficient(1, f, x, weights)
+
+        eq = Eq(f.forward.dx + f.backward.dx, 0, coefficients=Substitutions(coeffs))
+
+        assert str(eq.evaluate.lhs) == expected