diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index bb81e712a9..0f94304f97 100644 --- a/devito/finite_differences/finite_difference.py +++ b/devito/finite_differences/finite_difference.py @@ -6,7 +6,7 @@ from devito.logger import warning -from .differentiable import DiffDerivative, EvalDerivative, Weights +from .differentiable import DiffDerivative, EvalDerivative, Weights, merge_growth from .tools import ( centered, check_input, direct, fd_weights_registry, generate_indices, left, process_weights, right, transpose @@ -236,8 +236,12 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici # With a `subdomain`, `expr` is treated as zero outside it: every stencil tap # is multiplied by the indicator of the point it reads, 1 if `subdomain` # spans the whole of `dim`. The derivative then extends past `subdomain` by - # the stencil radius - growth = {dim.root: indices.radius} if subdomain is not None else {} + # the stencil radius, on top of the growth of `expr` itself (e.g., the inner + # derivative of a cross derivative) + if subdomain is not None: + growth = merge_growth([expr._growth, {dim.root: indices.radius}]) + else: + growth = {} if not expand and indices.expr is not None: weights = Weights(name='w', dimensions=indices.free_dim, diff --git a/devito/finite_differences/tools.py b/devito/finite_differences/tools.py index 4d816866fb..f517e4fd25 100644 --- a/devito/finite_differences/tools.py +++ b/devito/finite_differences/tools.py @@ -299,6 +299,10 @@ def generate_indices(expr, dim, order, side=None, matvec=None, x0=None, nweights iexpr = x0 + d * dim.spacing return IndexSet(dim, expr=iexpr), x0 + # Normalize x0 `Dimension` to match index if the latter is a `SubDimension` of + # the former + x0 = x0.subs({d.root: d for d in expr.dimensions if d.root is dim and d.is_Sub}) + # Evaluation point relative to the expression's grid mid = (x0 - expr.indices_ref[dim]).subs({dim: 0, dim.spacing: 1}) diff --git a/devito/types/dense.py b/devito/types/dense.py index 30f543a47c..fbfcd2dac2 100644 --- a/devito/types/dense.py +++ b/devito/types/dense.py @@ -1165,16 +1165,18 @@ def __staggered_setup__(cls, dimensions, staggered=None, **kwargs): processed = staggered else: # Staggering is not NODE or CELL or None - # therefore it's a tuple of dimensions + # therefore it's a tuple of dimensions. These may be the roots of + # `dimensions` (e.g., `staggered=x` for a Function on a SubDomain) + staggered = set(as_tuple(staggered)) processed = [] for d in dimensions: - if d in as_tuple(staggered): + if {d, d.root} & staggered: processed.append(sympy.S.One) - elif -d in as_tuple(staggered): + elif {-d, -d.root} & staggered: processed.append(sympy.S.NegativeOne) else: processed.append(sympy.S.Zero) - return Staggering(*processed, getters=dimensions) + return Staggering(*processed, getters=[d.root for d in dimensions]) @classmethod def __indices_setup__(cls, *args, **kwargs): diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index ee93c00488..f15dfe78b0 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -1861,7 +1861,7 @@ def test_adjoint_unexpanded_3d(self): assert self.dot_test(grid, subdomain, 'dy', 8, opt=opt) < 1e-12 - @pytest.mark.parametrize('deriv', ['dx', 'dy']) + @pytest.mark.parametrize('deriv', ['dx', 'dy', 'dxdy']) def test_adjoint_2d(self, deriv): grid = Grid(shape=(20, 22), extent=(19., 21.), dtype=np.float64) subdomain = Region('s', {'x': ('left', 6), 'y': ('middle', 4, 5)}, grid=grid) diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index eff6758636..bcde422cba 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -6,8 +6,8 @@ from conftest import assert_structure, opts_tiling from devito import ( - Border, Buffer, ConditionalDimension, Constant, Eq, Function, Grid, Lt, Operator, - SparseFunction, SparseTimeFunction, SubDimension, SubDomain, SubDomainSet, + Border, Buffer, ConditionalDimension, Constant, Derivative, Eq, Function, Grid, Lt, + Operator, SparseFunction, SparseTimeFunction, SubDimension, SubDomain, SubDomainSet, TensorFunction, TimeFunction, VectorFunction, solve ) from devito.exceptions import InvalidArgument @@ -1601,6 +1601,40 @@ def define(self, dimensions): assert np.all(np.isclose(fdx.data[:], gdx.data[2:-2, 3:-1])) assert np.all(np.isclose(fdy.data[:], gdy.data[2:-2, 3:-1])) + @pytest.mark.parametrize('s_o', [2, 4]) + @pytest.mark.parametrize('axis', [0, 1]) + def test_derivatives_staggered(self, s_o, axis): + """ + Derivatives of Functions on SubDomains can be evaluated at points staggered + along the derivative Dimension + """ + + class Middle(SubDomain): + + name = 'middle' + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 2, 2), y: ('middle', 3, 1)} + + grid = Grid(shape=(10, 10), extent=(9., 9.)) + d = grid.dimensions[axis] + mid = Middle(grid=grid) + + f = Function(name='f', grid=mid, space_order=s_o) + g = Function(name='g', grid=grid, space_order=s_o) + fd, gd = (Function(name=i, grid=grid, staggered=d) for i in ('fd', 'gd')) + + f.data[:] = np.random.default_rng(0).random(f.shape) + g.data[2:-2, 3:-1] = f.data + + Operator([Eq(fd, Derivative(f, d), subdomain=mid), + Eq(gd, Derivative(g, d), subdomain=mid)])() + + # Same stencil, so bitwise identical + assert np.any(fd.data != 0) + assert np.all(fd.data == gd.data) + def test_operator_shortcuts(self): """ Test that shortcuts for Laplacian, divergence, etc all work with Functions @@ -1982,14 +2016,14 @@ def define(self, dimensions): return {x: ('left', 2), y: y, z: z} sd1 = SD1(grid=grid) - f = TimeFunction(name='f', grid=grid, time_order=1, space_order=2, - staggered=(x,)) + f = TimeFunction(name='f', grid=grid, time_order=1, space_order=2) g = TimeFunction(name='g', grid=sd1, time_order=1, space_order=2, staggered=(x,)) eq = Eq(g, g + f.dx) eqe = eq.evaluate - assert eqe.rhs == g + f.dx(x0=x).evaluate._subs(x, g.dimensions[1]) + x0 = x + x.spacing/2 + assert eqe.rhs == g + f.dx(x0=x0).evaluate._subs(x, g.dimensions[1]) class TestSubDomainArguments: