diff --git a/devito/finite_differences/derivative.py b/devito/finite_differences/derivative.py index 08feb17839..de28b50859 100644 --- a/devito/finite_differences/derivative.py +++ b/devito/finite_differences/derivative.py @@ -14,7 +14,7 @@ from devito.warnings import warn from .differentiable import Add, Differentiable, Mul, diffify, interp_for_fd -from .finite_difference import cross_derivative, generic_derivative +from .finite_difference import cross_derivative, generic_derivative, indices_at from .rsfd import d45 from .tools import direct, transpose @@ -575,6 +575,12 @@ def _eval_fd(self, expr, **kwargs): shited derivative. - 4: Apply substitutions. """ + # Differentiation is linear, and a sum of terms at different staggered + # locations must use it: `Add` reports its first argument's location, + # so `x0` would shift the other terms off the point they sat at. + if expr.is_Add and any(len(indices_at(expr, d)) > 1 for d in self.dims): + return expr.func(*[self._eval_fd(a, **kwargs) for a in expr.args]) + # Step 1: Evaluate non-derivative x0. We currently enforce a simple 2nd order # interpolation to avoid very expensive finite differences on top of it x0_deriv = self._filter_dims(self.x0) diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index 11c2946ed1..de2e92898d 100644 --- a/devito/finite_differences/finite_difference.py +++ b/devito/finite_differences/finite_difference.py @@ -100,6 +100,29 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, side=None, **kw return expr +def indices_at(expr, dim): + """ + The locations `expr`'s terms sit at along `dim`. + + Terms with no location of their own, a scalar say, contribute none. + """ + indices = set() + for i in (expr.args if expr.is_Add else (expr,)): + try: + indices.add(i.indices_ref[dim]) + except (AttributeError, KeyError, IndexError, TypeError): + continue + return indices + + +def index_at(expr, dim): + """ + Where `expr` sits along `dim`, or None if it does not say. + """ + indices = indices_at(expr, dim) + return indices.pop() if len(indices) == 1 else None + + @check_input def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None, coefficients='taylor', expand=True, weights=None, side=None): @@ -139,8 +162,9 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None, if deriv_order == 1 and fd_order == 2 and side is None: fd_order = 1 - # Zeroth order derivative is just the expression itself if not shifted - if deriv_order == 0 and not x0: + # Zeroth order is the identity when `expr` already sits at `x0`, not a + # stencil centred there. + if deriv_order == 0 and (not x0 or index_at(expr, dim) == x0.get(dim)): return expr # Enforce stable time coefficients diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 81b12fa53b..bb570dbc2d 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -17,8 +17,8 @@ from devito.exceptions import CompilationError from devito.ir.cgen.printer import get_printer from devito.ir.iet.nodes import ( - BlankLine, Call, Expression, ExpressionBundle, Iteration, Lambda, ListMajor, Node, - Section, _same_as_before + BlankLine, Call, Definition, Expression, ExpressionBundle, Iteration, Lambda, + ListMajor, Node, Section, _same_as_before ) from devito.ir.support.space import Backward from devito.symbolics import ( @@ -1275,6 +1275,16 @@ def visit_Call(self, o: Call, **kwargs) -> Iterator[ApplicationType]: except (AttributeError, TypeError): yield from self._visit(i) + def visit_Definition(self, o: Definition, **kwargs) -> Iterator[ApplicationType]: + # A LocalObject carries expressions in its constructor arguments and in + # its initializer, both of which end up in the generated code + f = o.function + for i in (*getattr(f, 'cargs', ()), getattr(f, 'initvalue', None)): + try: + yield from i.find(self.match) + except (AttributeError, TypeError): + continue + class IsPerfectIteration(Visitor): diff --git a/devito/types/utils.py b/devito/types/utils.py index 8c5617f125..ecb2b360f6 100644 --- a/devito/types/utils.py +++ b/devito/types/utils.py @@ -44,17 +44,21 @@ class Stagger(Tag): class DimensionTuple(EnrichedTuple): - def __getitem_hook__(self, dim): + def _getter(self, dim): + # Exact hit first: a derived Dimension carries its parent in + # `_defines`, so an overlap test alone matches the parent's entry. + if dim in self.getters: + return dim for d in self.getters: if d._defines & dim._defines: - return self.getters[d] + return d raise KeyError + def __getitem_hook__(self, dim): + return self.getters[self._getter(dim)] + def dindex(self, dim): - for d in self.getters: - if d._defines & dim._defines: - return list(self.getters).index(d) - raise KeyError + return list(self.getters).index(self._getter(dim)) class Staggering(DimensionTuple): diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index 7ed57b6c91..362d3d1201 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -1461,3 +1461,43 @@ def test_unevaluated(self): assert Derivative(self.x, self.t) assert Derivative(self.x, self.y, self.t) assert Derivative(self.x, (self.x, 0)) + + +@pytest.mark.parametrize('expand', [True, False]) +@pytest.mark.parametrize('deriv_order', [0, 1, 2]) +def test_deriv_sum_mixed_staggering(expand, deriv_order): + """ + A shifted derivative is linear: `D(a + b) == D(a) + D(b)`, at every order. + + Broke for terms at different staggered locations, `Add` reporting only its + first argument's. + """ + so = 8 + grid = Grid(shape=(41, 41), extent=(40., 40.)) + x, y = grid.dimensions + + vx = Function(name='vx', grid=grid, space_order=so, staggered=x) + vy = Function(name='vy', grid=grid, space_order=so, staggered=y) + out = Function(name='out', grid=grid, space_order=so, staggered=(x, y)) + + rng = np.random.default_rng(3) + for f in (vx, vy): + f.data[:] = rng.normal(size=f.shape) + + def shifted(expr): + return expr.diff(y, deriv_order=deriv_order, fd_order=2, + x0={y: y + y.spacing/2}) + + def run(expr): + out.data[:] = 0. + Operator(Eq(out, expr), opt=('advanced', {'expand': expand})).apply() + return np.array(out.data) + + s = slice(so + 3, -(so + 3)) + together = run(shifted(vx.dy + vy.dx))[s, s] + apart = (run(shifted(vx.dy)) + run(shifted(vy.dx)))[s, s] + + assert np.linalg.norm(apart) > 0 + # float32 reassociation only: the two forms sum the same terms in a + # different order + assert np.linalg.norm(together - apart) / np.linalg.norm(apart) < 1e-5 diff --git a/tests/test_linearize.py b/tests/test_linearize.py index c2b424d934..c78880924d 100644 --- a/tests/test_linearize.py +++ b/tests/test_linearize.py @@ -8,7 +8,7 @@ ) from devito.ir import Call, Callable, DummyExpr, Expression, FindNodes, SymbolRegistry from devito.passes import Graph, generate_macros, linearize -from devito.types import Array, Bundle, DefaultDimension +from devito.types import Array, Bundle, CustomDimension, DefaultDimension def test_basic(): @@ -716,3 +716,24 @@ def test_cire_n_strides(): # NOTE: not exact equality because `op2` slightly changes the order of # arithmetic operations, which in turn causes some rounding differences assert np.allclose(u.data, u1.data, rtol=1e-4) + + +def test_bundle_derived_dim_stride(): + """ + A Bundle's stride comes from the Dimension asked for, not from its parent. + + `rp._defines` contains `p`, so an overlap lookup gave the Bundle the number + of points as innermost stride instead of the number of weights. + """ + grid = Grid(shape=(4, 4)) + p = DefaultDimension(name='p', default_value=5) + rp = CustomDimension(name='rp', parent=p, symbolic_size=2) + + w0 = Function(name='w0', dimensions=(p, rp), shape=(5, 2)) + w1 = Function(name='w1', dimensions=(p, rp), shape=(5, 2)) + bundle = Bundle(name='w0w1', components=(w0, w1), grid=grid) + + assert rp._defines & p._defines # the overlap that used to mislead + assert bundle.symbolic_shape[p] is not bundle.symbolic_shape[rp] + assert bundle.symbolic_shape[rp] == 2 + assert bundle.symbolic_shape.dindex(rp) == 1 diff --git a/tests/test_visitors.py b/tests/test_visitors.py index b5d12f81d5..eb925cb75f 100644 --- a/tests/test_visitors.py +++ b/tests/test_visitors.py @@ -1,3 +1,5 @@ +from ctypes import c_void_p + import cgen as c import pytest from sympy import Mod @@ -5,11 +7,12 @@ from devito import Eq, Function, Grid, Min, Operator, TimeFunction, sin from devito.ir.equations import DummyEq from devito.ir.iet import ( - Block, Call, Callable, Conditional, Expression, FindApplications, FindNodes, - FindSections, FindSymbols, FindWithin, IsPerfectIteration, Iteration, MapNodes, - Transformer, Uxreplace, printAST + Block, Call, Callable, Conditional, Definition, Expression, FindApplications, + FindNodes, FindSections, FindSymbols, FindWithin, IsPerfectIteration, Iteration, + MapNodes, Transformer, Uxreplace, printAST ) -from devito.types import Array, SpaceDimension, Symbol +from devito.symbolics import ListInitializer +from devito.types import Array, LocalObject, SpaceDimension, Symbol @pytest.fixture(scope="module") @@ -422,3 +425,25 @@ def test_find_apps_nested_calls(exprs, iters): block = iters[0](iters[1](exprs + [call])) assert len(FindApplications().visit(block)) == 1 + + +def test_find_apps_in_definition(): + """ + A LocalObject carries expressions in its constructor arguments and in its + initializer, and both end up in the generated code, so both must be + searched -- otherwise e.g. `generate_macros` would leave the macros they + apply undefined. + """ + s = Symbol(name='s') + + class DummyObject(LocalObject): + dtype = c_void_p + + obj = DummyObject(name='obj', initvalue=ListInitializer([Min(s, 1)])) + assert FindApplications().visit(Definition(obj)) == {Min(s, 1)} + + obj = DummyObject(name='obj', cargs=(Min(s, 2),)) + assert FindApplications().visit(Definition(obj)) == {Min(s, 2)} + + # A LocalObject with neither must not trip the search up + assert FindApplications().visit(Definition(DummyObject(name='obj'))) == set()