diff --git a/devito/finite_differences/coefficients.py b/devito/finite_differences/coefficients.py index 4f7e48cf3a..c686d4a997 100644 --- a/devito/finite_differences/coefficients.py +++ b/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} diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index a472ae6e02..3c2b540e14 100644 --- a/devito/finite_differences/finite_difference.py +++ b/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) diff --git a/devito/finite_differences/tools.py b/devito/finite_differences/tools.py index e97436c482..11cadb7a3d 100644 --- a/devito/finite_differences/tools.py +++ b/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,6 +165,8 @@ 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 ------- @@ -172,7 +174,8 @@ def generate_indices(func, dim, order, side=None, x0=None): """ # 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)] diff --git a/devito/types/equation.py b/devito/types/equation.py index 4b86a59ec9..c53067ccf5 100644 --- a/devito/types/equation.py +++ b/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) diff --git a/tests/test_symbolic_coefficients.py b/tests/test_symbolic_coefficients.py index fa8a6d2617..1982486f05 100644 --- a/tests/test_symbolic_coefficients.py +++ b/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