From a707eb5e5bc1382a694f88fd245e79cb930e4ce6 Mon Sep 17 00:00:00 2001 From: mloubout Date: Mon, 28 Sep 2026 21:04:36 -0400 Subject: [PATCH 1/3] compiler: Shift SubDomain Functions in finite differences of mixed expressions A derivative of an expression mixing Grid Functions, e.g. f(x), with Functions defined on a SubDomain, e.g. p(ix), only substituted the derivative Dimension in the stencil. Functions indexed by the sibling Dimension stayed unshifted, so (f + p).dx silently dropped p and (f*p).dx never shifted p. Differentiable.shift now shifts every space Dimension sharing the root of the given Dimension, and make_derivative builds the stencil points with it, in both the expanded and unexpanded paths. Tests cover Add and Mul with both expansion modes, and a one-face CPML forward/adjoint dot test whose adjoint term mixes Grid and SubDomain Functions inside .dx.T. --- devito/finite_differences/differentiable.py | 12 +- .../finite_differences/finite_difference.py | 4 +- tests/test_derivatives.py | 109 +++++++++++++++++- 3 files changed, 119 insertions(+), 6 deletions(-) diff --git a/devito/finite_differences/differentiable.py b/devito/finite_differences/differentiable.py index 1e395e1a38..2edf20a3ea 100644 --- a/devito/finite_differences/differentiable.py +++ b/devito/finite_differences/differentiable.py @@ -329,8 +329,18 @@ def shift(self, dim, shift): """ Shift expression by `shift` along the Dimension `dim`. For example u.shift(x, x.spacing) = u(x + h_x). + + Every other space Dimension of the expression sharing the root of `dim` + is shifted too: an expression may mix Functions defined on the Grid, + e.g. `f(x)`, with Functions defined on a SubDomain, e.g. `p(ix)`. """ - return self._subs(dim, dim + shift) + from devito.symbolics import retrieve_dimensions # noqa + + expr = self + for d in retrieve_dimensions(self, mode='unique'): + if d is not dim and dim.root in d._defines and not d.is_NonlinearDerived: + expr = expr._subs(d, d + shift) + return expr._subs(dim, dim + shift) @property def laplace(self): diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index 73becb8b84..5cb5d1a1c0 100644 --- a/devito/finite_differences/finite_difference.py +++ b/devito/finite_differences/finite_difference.py @@ -228,7 +228,7 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici # Inject the StencilDimension # E.g. `x + i*h_x` into `f(x)` s.t. `f(x + i*h_x)` - expr = expr._subs(dim, indices.expr) + expr = expr.shift(dim, indices.expr - dim) # Re-evaluate any off-the-grid Functions potentially impacted by the FD # unless a pure number @@ -242,7 +242,7 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici terms = [] for i, c in zip(indices, weights, strict=True): # The FD term - term = expr._subs(dim, i) * c + term = expr.shift(dim, i - dim) * c # Re-evaluate any off-the-grid Functions potentially impacted by the FD # unless a pure number diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index b5169efbcd..46fad9cdeb 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -4,9 +4,9 @@ from conftest import assert_structure from devito import ( - NODE, ConditionalDimension, Eq, Function, Grid, Operator, TensorFunction, - TensorTimeFunction, TimeFunction, VectorFunction, centered, cos, curl, div, grad, - laplace, left, right, sin + NODE, ConditionalDimension, Eq, Function, Grid, Inc, Operator, SubDomain, + TensorFunction, TensorTimeFunction, TimeFunction, VectorFunction, centered, cos, curl, + div, grad, laplace, left, right, sin ) from devito.finite_differences import Derivative, Differentiable, diffify from devito.finite_differences.differentiable import ( @@ -1566,3 +1566,106 @@ def run(expr): # 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 + + +class Region(SubDomain): + """SubDomain given by its `define` mapping for each Dimension name.""" + + def __init__(self, name, regions, **kwargs): + self.name = name + self.regions = regions + super().__init__(**kwargs) + + def define(self, dimensions): + return {d: self.regions.get(d.name, d) for d in dimensions} + + +@pytest.mark.parametrize('expand', [True, False]) +@pytest.mark.parametrize('op', ['add', 'mul']) +def test_deriv_mixed_subdomain_functions(op, expand): + """ + A derivative shifts every Function of the expression, including those + defined on a SubDomain, whose Dimension (`ix`) differs from the Grid's (`x`). + + Used to leave `p(ix)` unshifted, so `(f + p).dx` dropped `p` altogether. + """ + so = 4 + grid = Grid(shape=(24,), extent=(23.,)) + strip = Region('strip', {'x': ('left', 8)}, grid=grid) + + f = Function(name='f', grid=grid, space_order=so) + p = Function(name='p', grid=strip, space_order=so) + # Reference: `p` extended by zero on the whole Grid + p_full = Function(name='p_full', grid=grid, space_order=so) + out = Function(name='out', grid=grid, space_order=so) + ref = Function(name='ref', grid=grid, space_order=so) + + rng = np.random.default_rng(0) + f.data[:] = rng.normal(size=f.shape) + p.data[:] = rng.normal(size=p.shape) + p_full.data[:8] = p.data + + combine = {'add': lambda u, v: u + v, 'mul': lambda u, v: u * v}[op] + Operator([Eq(out, combine(f, p).dx, subdomain=strip), + Eq(ref, combine(f, p_full).dx, subdomain=strip)], + opt=('advanced', {'expand': expand})).apply() + + assert np.linalg.norm(ref.data) > 0 + assert np.allclose(out.data, ref.data, rtol=1e-6, atol=1e-6) + + +def test_deriv_mixed_subdomain_cpml_adjoint(): + """ + Dot test of a one-face CPML derivative whose memory variable lives on the + CPML layer only. + + forward: out = D q + psi[n+1] (psi on the layer) + psi[n+1] = b psi[n] + a D q + adjoint: q_bar = D^T (out_bar + a (psi_bar + out_bar)) + + The adjoint CPML term mixes Grid and SubDomain Functions inside `.dx.T`, and + is nonzero on the layer grown by the stencil radius. `a` vanishes outside + the layer, so no zero-padded work field is needed. + """ + nx, width, so, nt = 32, 6, 8, 8 + grid = Grid(shape=(nx,), extent=(float(nx - 1),), dtype=np.float64) + time = grid.time_dim + x = grid.dimensions[0] + cpml = Region('cpml', {'x': ('left', width)}, grid=grid) + cpml_r = Region('cpml_r', {'x': ('left', width + so // 2)}, grid=grid) + + kwargs = {'grid': grid, 'save': nt, 'space_order': so, 'dtype': np.float64} + q = TimeFunction(name='q', **kwargs) + out = TimeFunction(name='out', **kwargs) + q_bar = TimeFunction(name='q_bar', **kwargs) + out_bar = TimeFunction(name='out_bar', **kwargs) + psi = TimeFunction(name='psi', grid=cpml, time_order=1, space_order=0, + dtype=np.float64) + # The adjoint loop over `cpml_r` reads `psi_bar` one more radius away + psi_bar = TimeFunction(name='psi_bar', grid=cpml, time_order=1, + space_order=(so, so, so), dtype=np.float64) + a = Function(name='a', grid=grid, dimensions=(x,), shape=(nx,), space_order=so, + dtype=np.float64) + b = Function(name='b', grid=grid, dimensions=(x,), shape=(nx,), dtype=np.float64) + a.data[:width] = -0.1 + b.data[:width] = 0.8 + + forward = Operator([Eq(out, q.dx), + Eq(psi.forward, b * psi + a * q.dx, subdomain=cpml), + Eq(out, out + psi.forward, subdomain=cpml)]) + value = psi_bar + out_bar + adjoint = Operator([ + Eq(psi_bar.backward, b * value, subdomain=cpml, implicit_dims=(time,)), + Eq(q_bar, out_bar.dx.T, implicit_dims=(time,)), + Inc(q_bar, (a * value).dx.T, subdomain=cpml_r, implicit_dims=(time,)) + ]) + + rng = np.random.default_rng(7) + q.data[:] = rng.normal(size=q.shape) + out_bar.data[:] = rng.normal(size=out_bar.shape) + forward.apply(time_m=0, time_M=nt - 1) + adjoint.apply(time_m=0, time_M=nt - 1) + + lhs = np.vdot(out.data, out_bar.data) + rhs = np.vdot(q.data, q_bar.data) + assert np.isclose(lhs, rhs, rtol=1e-12) From 5f2b2aad35c8b0ada8e01b639f2281a9ddc3c219 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 29 Sep 2026 16:17:03 -0400 Subject: [PATCH 2/3] compiler: Do not alias untranslatable StencilDimensions in CIRE CIRE builds aliases by translating Indexeds. An unbound StencilDimension appearing outside of any Indexed, e.g. `i0` in `x + i0` within the body of an IndexDerivative, is not translated, so the alias would be evaluated at the wrong stencil point. For example, the invariants pass hoisted such an expression into an array over `x`, dropping `i0`. `_do_generate` now leaves the terms reading such StencilDimensions out of the alias, or discards the candidate altogether. --- devito/passes/clusters/aliases.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/devito/passes/clusters/aliases.py b/devito/passes/clusters/aliases.py index 22f47b5c94..bc757be998 100644 --- a/devito/passes/clusters/aliases.py +++ b/devito/passes/clusters/aliases.py @@ -16,8 +16,8 @@ from devito.passes.clusters.cse import _cse from devito.passes.clusters.utils import expose_tuning_knobs from devito.symbolics import ( - Uxmapper, estimate_cost, retrieve_functions, reuse_if_untouched, search, sympy_dtype, - uxreplace + Uxmapper, estimate_cost, retrieve_dimensions, retrieve_functions, reuse_if_untouched, + search, sympy_dtype, uxreplace ) from devito.tools import ( Reconstructable, Stamp, as_mapper, as_tuple, flatten, generator, is_integer, split, @@ -282,6 +282,18 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None): terms = cbk_compose(i) + # Aliases only translate Indexeds, so a term reading an unbound + # StencilDimension outside of them would be evaluated at the + # wrong stencil point: such terms are left out of the alias. + # E.g., the `halo=0` indicator `MAX(0, MIN(1, x + i0 - x_m + 1))` + # in `u[x + i0]*w[i0]*MAX(0, MIN(1, x + i0 - x_m + 1))` + if terms: + terms = [t for t in terms if not untranslatable(t)] + if not terms: + continue + elif untranslatable(i): + continue + # Make sure we won't break any data dependencies if terms: free_symbols = set().union(*[i.free_symbols for i in terms]) @@ -1612,6 +1624,16 @@ def split_coeff(expr): return maybe_coeff, others +def untranslatable(expr): + """ + The unbound StencilDimensions of `expr` appearing outside of any Indexed, + which translating `expr` into an alias would not shift. `retrieve_dimensions` + does not look within Indexeds, so the intersection with `unbounded(expr)` + leaves out those appearing only in Indexeds, which are translated. + """ + return unbounded(expr) & retrieve_dimensions(expr, mode='unique') + + def nredundants(ispace, expr): """ The number of redundant Dimensions in `ispace` for `expr`. A Dimension is From 026764dfd78ae39b7089196bc565871b3592648e Mon Sep 17 00:00:00 2001 From: mloubout Date: Sun, 4 Oct 2026 15:48:13 -0400 Subject: [PATCH 3/3] api: Add Derivative halo=0 for SubDomain-restricted derivatives and adjoints `expr.dx(halo=0)` treats `expr` as zero outside the SubDomain of its equation, and `.T` keeps the flag. Hence `Inc(q_bar, out_bar.dx(halo=0).T, subdomain=S)` is the adjoint of `Eq(out, q.dx, subdomain=S)`, i.e. D^T R^T with R the restriction to S, without a user-side zero-padded work field. Equations without `halo=0` derivatives, or without a SubDomain, are evaluated as before. - An equation evaluates its rhs with `_eval_at(lhs, subdomain=S)`. A `halo=0` derivative records S, and multiplies each stencil tap by a branch-free integer MIN/MAX indicator of the grid point it reads being in S (SubDomain.indicator), which invariant hoisting computes once. Along the Dimensions it does not differentiate, it restricts its argument to S at the evaluation point. In a sum with such derivatives, the other terms are restricted to S at the evaluation point (SubDomain.restrict); factors are not. - The evaluated derivative records its stencil radius as its extent past S (Differentiable.extent), and the equation iterates over S grown by the largest one (SubDomain.grow, a SubDomain with a parent). Tensor equations aggregate their components. - It works for expanded, unexpanded and staggered derivatives. Derivatives of `halo=0` derivatives, MultiSubDomains and methods other than FD raise NotImplementedError. Tests: dot tests of `Eq(out, g.dx, subdomain=S)` against `Inc(g_bar, out_bar.dx(halo=0).T, subdomain=S)` for left/right/middle SubDomains, first and second derivatives, orders 2/4/8, Eq and Inc, 2D and unexpanded forms, in 1D and 3D; staggered dot tests against `-out_bar.dx(halo=0)`, as in elastic adjoints, including a staggered vector equation; sums, factors and vector equations; SubDomain growth; a forward check; and a one-face CPML operator. --- devito/finite_differences/derivative.py | 116 ++++- devito/finite_differences/differentiable.py | 93 +++- .../finite_differences/finite_difference.py | 26 +- devito/finite_differences/tools.py | 14 +- devito/passes/clusters/aliases.py | 7 +- devito/types/equation.py | 12 +- devito/types/grid.py | 107 ++++- devito/types/tensor.py | 13 +- tests/test_derivatives.py | 405 ++++++++++++++++++ 9 files changed, 751 insertions(+), 42 deletions(-) diff --git a/devito/finite_differences/derivative.py b/devito/finite_differences/derivative.py index de28b50859..bd7d8c4c16 100644 --- a/devito/finite_differences/derivative.py +++ b/devito/finite_differences/derivative.py @@ -102,7 +102,7 @@ def _fd_priority(self): __rargs__ = ('expr', '*dims') __rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs', - 'x0', 'method', 'weights') + 'x0', 'method', 'weights', 'halo', 'subdomain') def __new__(cls, expr, *dims, **kwargs): # Validate the input arguments `expr`, `dims` and `deriv_order` @@ -161,6 +161,8 @@ def __new__(cls, expr, *dims, **kwargs): obj._transpose = kwargs.get("transpose", direct) obj._method = kwargs.get("method", 'FD') obj._weights = cls._process_weights(**kwargs) + obj._halo = cls._validate_halo(kwargs.get("halo")) + obj._subdomain = kwargs.get("subdomain") ppsubs = kwargs.get("subs", kwargs.get("_ppsubs", [])) processed = [] @@ -177,6 +179,16 @@ def __new__(cls, expr, *dims, **kwargs): return obj + @staticmethod + def _validate_halo(halo): + """ + Validate `halo`. Only None (read the argument everywhere) and 0 (treat + the argument as zero outside the equation's SubDomain) are supported. + """ + if halo not in (None, 0): + raise ValueError(f"Expected halo=None or halo=0, got halo={halo}") + return halo + @staticmethod def _validate_expr(expr): """ @@ -325,6 +337,8 @@ def _process_weights(cls, **kwargs): def __call__(self, x0=None, fd_order=None, side=None, method=None, **kwargs): weights = kwargs.get('weights', kwargs.get('w')) rkw = {} + if 'halo' in kwargs: + rkw['halo'] = kwargs['halo'] if side is not None: rkw['side'] = side if method is not None: @@ -457,6 +471,26 @@ def side(self): def transpose(self): return self._transpose + @property + def halo(self): + """ + None if the argument is read everywhere, 0 if it is treated as zero + outside the SubDomain of the equation the Derivative belongs to. + """ + return self._halo + + @property + def subdomain(self): + """ + With halo=0, the SubDomain outside of which the argument is treated as + zero, set upon evaluation within an equation restricted to it. + """ + return self._subdomain + + @cached_property + def _has_zero_halo(self): + return self.halo is not None or self.expr._has_zero_halo + @property def is_TimeDependent(self): return self.expr.is_TimeDependent @@ -481,26 +515,21 @@ def T(self): return self._rebuild(transpose=adjoint) - def _eval_at(self, func, interp_mode='direct', **kwargs): + def _eval_at(self, func, interp_mode='direct', subdomain=None, **kwargs): """ Evaluates the derivative at the location of `func`. It is necessary for staggered setup where one could have Eq(u(x + h_x/2), v(x).dx)) in which case v(x).dx has to be computed at x=x + h_x/2. + + With halo=0, the argument is treated as zero outside `subdomain`, which + the Derivative records. """ - # No staggering, don't waste time - if not self.expr.staggered and not func.staggered: - return self - # If an x0 already exists or evaluating at the same function (i.e u = u.dx) - # do not overwrite it - if self.x0 or self.side is not None or func.function is self.expr.function: - return self - # For basic equation of the form f = Derivative(g, ...) we can just - # compare staggering - if self.expr.staggered == func.staggered and self.expr.is_Function: - return self - # Time derivatives are not affected by space staggering - if all(d.is_Time for d in self.dims): - return self + rkw = {} + if subdomain is not None and self.halo is not None: + rkw['subdomain'] = subdomain + + if not self._is_relocated(func): + return self._rebuild(**rkw) if rkw else self # Check if x0's keys come from a DerivedDimension x0 = func.indices_ref.getters @@ -519,7 +548,7 @@ def _eval_at(self, func, interp_mode='direct', **kwargs): # e.g f.dx(x0={x: x + h_x/2}).subs({x: ix}) psubs[sd] = d nx0[sd] = nx0.pop(d)._subs(d, sd) - rkw = {'x0': nx0} + rkw['x0'] = nx0 if psubs: rkw['subs'] = (psubs,) @@ -534,7 +563,8 @@ def _eval_at(self, func, interp_mode='direct', **kwargs): return self._rebuild(self.expr, **rkw) args = [self.expr.func(*v) for v in mapper.values()] args.extend([a for a in self.expr.args if a not in self.expr._args_diff]) - args = [self._rebuild(a)._eval_at(func, interp_mode=interp_mode, **kwargs) + args = [self._rebuild(a)._eval_at(func, interp_mode=interp_mode, + subdomain=subdomain, **kwargs) for a in args] return self.expr.func(*args) elif self.expr.is_Mul: @@ -549,6 +579,25 @@ def _eval_at(self, func, interp_mode='direct', **kwargs): # the expression as is. return self._rebuild(self.expr, **rkw) + def _is_relocated(self, func): + """ + True if the Derivative must be evaluated at the location of `func`, e.g. + with `func` and the argument staggered apart. + """ + # No staggering, don't waste time + if not self.expr.staggered and not func.staggered: + return False + # If an x0 already exists or evaluating at the same function (i.e u = u.dx) + # do not overwrite it + if self.x0 or self.side is not None or func.function is self.expr.function: + return False + # For basic equation of the form f = Derivative(g, ...) we can just + # compare staggering + if self.expr.staggered == func.staggered and self.expr.is_Function: + return False + # Time derivatives are not affected by space staggering + return not all(d.is_Time for d in self.dims) + def _evaluate(self, **kwargs): # Evaluate finite-difference. # NOTE: `evaluate` and `_eval_fd` split for potential future different @@ -581,6 +630,13 @@ def _eval_fd(self, expr, **kwargs): 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]) + # The SubDomain mask read by a halo=0 derivative can't be shifted again + if expr._has_zero_halo: + raise NotImplementedError( + f"{self} differentiates a derivative with halo=0, which is not " + "supported" + ) + # 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) @@ -598,6 +654,15 @@ def _eval_fd(self, expr, **kwargs): # otherwise an IndexSum will returned expand = kwargs.get('expand', True) + # With halo=0, `expr` is treated as zero outside the equation's SubDomain + subdomain = self.subdomain + if subdomain is not None and (subdomain.is_MultiSubDomain or + self.method != 'FD'): + raise NotImplementedError( + f"halo=0 is only supported with method='FD' on a SubDomain, not " + f"with method={self.method} on {subdomain}" + ) + # Step 3: Evaluate FD of the new expression if self.method == 'RSFD': assert len(self.dims) == 1 @@ -607,18 +672,29 @@ def _eval_fd(self, expr, **kwargs): assert self.method == 'FD' res = cross_derivative(expr, self.dims, self.fd_order, self.deriv_order, matvec=self.transpose, x0=x0_deriv, expand=expand, - side=self.side, weights=self.weights) + side=self.side, weights=self.weights, + subdomain=subdomain) else: assert self.method == 'FD' res = generic_derivative(expr, self.dims[0], self.fd_order[0], self.deriv_order[0], weights=self.weights, side=self.side, matvec=self.transpose, - x0=self.x0, expand=expand) + x0=self.x0, expand=expand, + subdomain=subdomain) # Step 4: Apply substitutions for e in self._ppsubs: res = res.xreplace(e) + # With `halo=0`, the stencil taps above are zeroed outside the SubDomain + # along `self.dims` only. Along any other Dimension, the argument is read + # at the evaluation point, which other `halo=0` derivatives may push + # outside the SubDomain. E.g., in `f.dx(halo=0) + f.dy(halo=0)`, the + # equation iterates over the SubDomain grown along both `x` and `y`, so + # `f.dx` must also be zeroed where `y` lies outside the SubDomain + if subdomain is not None: + res = subdomain._restrict(res, exclude={d.root for d in self.dims}) + return res def _eval_expand_nest(self, **hints): diff --git a/devito/finite_differences/differentiable.py b/devito/finite_differences/differentiable.py index 2edf20a3ea..2b8b6d5080 100644 --- a/devito/finite_differences/differentiable.py +++ b/devito/finite_differences/differentiable.py @@ -191,6 +191,26 @@ def _eval_at(self, func, **kwargs): for a in self.args # false positive: lambda is invoked in-place ]) + @cached_property + def _has_zero_halo(self): + """True if the expression has derivatives with `halo=0`.""" + return any(a._has_zero_halo for a in self._args_diff) + + @cached_property + def _growth(self): + """ + Number of points, per root Dimension, by which the evaluated expression + extends past the SubDomain that its `halo=0` derivatives restrict their + argument to: the largest stencil radius of such derivatives. Empty if the + expression has none. + + For example, with `S` a SubDomain restricting `x` and an 8th-order `g`, + `g.dx(halo=0)` evaluated in an equation on `S` has a growth of `{x: 4}`: + the equation must iterate over `S` grown by 4 points along `x` (see + `DerivedSubDomain`). + """ + return merge_growth(a._growth for a in self._args_diff) + def _subs(self, old, new, **hints): if old == self: return new @@ -540,6 +560,18 @@ def highest_priority(diff_op, candidates=None): return prio_func +def merge_growth(growths): + """ + Merge the growths `growths`, each a mapping from root Dimension to a number + of points, keeping the largest per Dimension (see `Differentiable._growth`). + """ + growth = {} + for i in growths: + for d, v in i.items(): + growth[d] = max(growth.get(d, 0), v) + return frozendict(growth) + + class DifferentiableOp(Differentiable): __sympy_class__ = None @@ -640,6 +672,19 @@ def __new__(cls, *args, **kwargs): return super().__new__(cls, *args, **kwargs) + def _eval_at(self, func, subdomain=None, **kwargs): + """ + Evaluate the sum at the location of `func`. + + The derivatives with `halo=0` extend the sum past `subdomain`, so the + other terms are restricted to `subdomain` at the evaluation point. + """ + expr = super()._eval_at(func, subdomain=subdomain, **kwargs) + if subdomain is None or not expr.is_Add or not expr._has_zero_halo: + return expr + halo = {a for a in expr._args_diff if a._has_zero_halo} + return self.func(*[a if a in halo else subdomain._restrict(a) for a in expr.args]) + class Mul(DifferentiableOp, sympy.Mul): __sympy_class__ = sympy.Mul @@ -1328,6 +1373,30 @@ def _subs(self, old, new, **hints): class DiffDerivative(IndexDerivative, DifferentiableOp): + """ + A Derivative evaluated in unexpanded form. + + Parameters + ---------- + growth : dict of {Dimension: int}, optional + For a derivative with `halo=0`, its stencil radius per root Dimension + (see `Differentiable._growth`). + *args, **kwargs + As for IndexDerivative. + """ + + __rkwargs__ = IndexDerivative.__rkwargs__ + ('growth',) + + def __new__(cls, *args, growth=None, **kwargs): + obj = super().__new__(cls, *args, **kwargs) + obj._growth = frozendict(growth or {}) + return obj + + @property + def growth(self): + """For a derivative with `halo=0`, its stencil radius per root Dimension.""" + return self._growth + def _eval_at(self, func, **kwargs): # Like EvalDerivative, a DiffDerivative must have already been evaluated # at a valid x0 and should not be re-evaluated at a different location @@ -1341,11 +1410,25 @@ def _eval_at(self, func, **kwargs): class EvalDerivative(DifferentiableOp, sympy.Add): + """ + A Derivative evaluated in expanded form, as the sum of its stencil taps. + + Parameters + ---------- + *args : expr-like + The stencil taps. + base : expr-like, optional + The expression the derivative was taken of. + growth : dict of {Dimension: int}, optional + For a derivative with `halo=0`, its stencil radius per root Dimension + (see `Differentiable._growth`). + """ + is_commutative = True - __rkwargs__ = ('base',) + __rkwargs__ = ('base', 'growth') - def __new__(cls, *args, base=None, **kwargs): + def __new__(cls, *args, base=None, growth=None, **kwargs): kwargs['evaluate'] = False # a+0 -> a @@ -1361,6 +1444,7 @@ def __new__(cls, *args, base=None, **kwargs): # In some rare cases (rebuild?) base may be obj itself base = base.base obj.base = base + obj._growth = frozendict(growth or {}) except AttributeError: # This might happen if e.g. one attempts a (re)construction with # one sole argument. The (re)constructed EvalDerivative degenerates @@ -1373,6 +1457,11 @@ def __new__(cls, *args, base=None, **kwargs): return obj + @property + def growth(self): + """For a derivative with `halo=0`, its stencil radius per root Dimension.""" + return self._growth + func = DifferentiableOp._rebuild # Since obj.base = base, then Differentiable.__eq__ leads to infinite recursion diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index 5cb5d1a1c0..29fa3b4c77 100644 --- a/devito/finite_differences/finite_difference.py +++ b/devito/finite_differences/finite_difference.py @@ -125,7 +125,8 @@ def index_at(expr, dim): @check_input def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None, - coefficients='taylor', expand=True, weights=None, side=None): + coefficients='taylor', expand=True, weights=None, side=None, + subdomain=None): """ Arbitrary-order derivative of a given expression. @@ -151,6 +152,10 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None, expand : bool, optional, default=True If True, the derivative is fully expanded as a sum of products, otherwise an IndexSum is returned. + subdomain : SubDomain, optional + If given, `expr` is treated as zero outside `subdomain` (see + `Derivative`'s `halo`): each stencil tap is multiplied by the + indicator of `subdomain` at the point it reads. Returns ------- @@ -171,7 +176,7 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None, coefficients = 'taylor' if dim.is_Time else expr.coefficients return make_derivative(expr, dim, fd_order, deriv_order, side, - matvec, x0, coefficients, expand, weights) + matvec, x0, coefficients, expand, weights, subdomain) # Backward compatibility @@ -180,7 +185,7 @@ def first_derivative(expr, dim, fd_order, **kwargs): def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coefficients, - expand, weights=None): + expand, weights=None, subdomain=None): # Always expand time derivatives to avoid issue with buffering and streaming. # Time derivative are almost always short stencils and won't benefit from # unexpansion in the rare case the derivative is not evaluated for time stepping. @@ -222,6 +227,12 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici if callable(expand): expand = expand(dim) + # 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 {} + if not expand and indices.expr is not None: weights = Weights(name='w', dimensions=indices.free_dim, initvalue=weights, dtype=expr.dtype) @@ -229,6 +240,8 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici # Inject the StencilDimension # E.g. `x + i*h_x` into `f(x)` s.t. `f(x + i*h_x)` expr = expr.shift(dim, indices.expr - dim) + if subdomain is not None: + expr = expr * subdomain._indicator(dim.root, indices.offset(indices.expr)) # Re-evaluate any off-the-grid Functions potentially impacted by the FD # unless a pure number @@ -236,13 +249,16 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici expr = expr._evaluate(expand=False) deriv = DiffDerivative( - expr*weights, {dim: indices.free_dim}, deriv_order=deriv_order + expr*weights, {dim: indices.free_dim}, deriv_order=deriv_order, + growth=growth ) else: terms = [] for i, c in zip(indices, weights, strict=True): # The FD term term = expr.shift(dim, i - dim) * c + if subdomain is not None: + term = term * subdomain._indicator(dim.root, indices.offset(i)) # Re-evaluate any off-the-grid Functions potentially impacted by the FD # unless a pure number @@ -250,6 +266,6 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici term = term.evaluate terms.append(term) - deriv = EvalDerivative(*terms, base=expr) + deriv = EvalDerivative(*terms, base=expr, growth=growth) return deriv diff --git a/devito/finite_differences/tools.py b/devito/finite_differences/tools.py index 8c9304b126..2e2d2847e1 100644 --- a/devito/finite_differences/tools.py +++ b/devito/finite_differences/tools.py @@ -1,4 +1,4 @@ -from functools import partial, wraps +from functools import cached_property, partial, wraps from itertools import product import numpy as np @@ -177,6 +177,18 @@ def __repr__(self): def spacing(self): return self.dim.spacing + def offset(self, index): + """ + Offset, in grid points, of the point `index` from `dim`: e.g. 2 for + `x + 2*h_x`, or `i0` for `x + i0*h_x` in unexpanded form. + """ + return sympify((index - self.dim) / self.spacing) + + @cached_property + def radius(self): + """Largest distance, in grid points, between `dim` and the points.""" + return max(abs(self.offset(i)) for i in self) + def transpose(self): """ Transpose the IndexSet. diff --git a/devito/passes/clusters/aliases.py b/devito/passes/clusters/aliases.py index bc757be998..4574321de3 100644 --- a/devito/passes/clusters/aliases.py +++ b/devito/passes/clusters/aliases.py @@ -282,9 +282,10 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None): terms = cbk_compose(i) - # Aliases only translate Indexeds, so a term reading an unbound - # StencilDimension outside of them would be evaluated at the - # wrong stencil point: such terms are left out of the alias. + # NOTE: a term reading an unbound StencilDimension outside of any + # Indexed is unsupported for now: aliases only translate Indexeds, + # so it would be evaluated at the wrong stencil point. It could in + # principle be hoisted too, but is currently left out of the alias. # E.g., the `halo=0` indicator `MAX(0, MIN(1, x + i0 - x_m + 1))` # in `u[x + i0]*w[i0]*MAX(0, MIN(1, x + i0 - x_m + 1))` if terms: diff --git a/devito/types/equation.py b/devito/types/equation.py index 42ddb07dae..e14e7f7b59 100644 --- a/devito/types/equation.py +++ b/devito/types/equation.py @@ -125,12 +125,20 @@ def _evaluate(self, **kwargs): """ if self._interp_mode is not None: kwargs['interp_mode'] = self._interp_mode + + subdomain = self.subdomain try: lhs = self.lhs._evaluate(**kwargs) - rhs = self.rhs._eval_at(self.lhs, **kwargs)._evaluate(**kwargs) + rhs = self.rhs._eval_at(self.lhs, subdomain=subdomain, **kwargs) + rhs = rhs._evaluate(**kwargs) except AttributeError: lhs, rhs = self._evaluate_args(**kwargs) - eq = self.func(lhs, rhs, subdomain=self.subdomain, + else: + # Derivatives with `halo=0` extend the rhs past the SubDomain + if subdomain is not None and self.rhs._has_zero_halo: + subdomain = subdomain._grow(rhs._growth) + + eq = self.func(lhs, rhs, subdomain=subdomain, coefficients=self.substitutions, implicit_dims=self._implicit_dims, interp_mode=self._interp_mode) diff --git a/devito/types/grid.py b/devito/types/grid.py index 6a5ec26c0e..378d2543b4 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -3,11 +3,12 @@ from itertools import product import numpy as np -from sympy import prod +import sympy from devito import configuration from devito.data import CENTER, LEFT, RIGHT from devito.deprecations import deprecations +from devito.finite_differences.elementary import Max, Min from devito.logger import warning from devito.mpi import MPI, Distributor, SubDistributor from devito.tools import ReducerMap, as_tuple, frozendict @@ -301,7 +302,7 @@ def interior(self): @property def volume_cell(self): """Volume of a single cell e.g h_x*h_y*h_z in 3D.""" - return prod(d.spacing for d in self.dimensions).subs(self.spacing_map) + return sympy.prod(d.spacing for d in self.dimensions).subs(self.spacing_map) @cached_property def spacing(self): @@ -649,12 +650,8 @@ def __subdomain_finalize_legacy__(self, grid): # Create the SubDomain's SubDimensions sub_dimensions = [] sdshape = [] - for k, v, s in zip( - self.define(grid.dimensions).keys(), - self.define(grid.dimensions).values(), - grid.shape, - strict=True - ): + regions = self.define(grid.dimensions) + for (k, v), s in zip(regions.items(), grid.shape, strict=True): if isinstance(v, Dimension): sub_dimensions.append(v) sdshape.append(s) @@ -725,6 +722,48 @@ def define(self, dimensions): """ raise NotImplementedError + def _grow(self, growth): + """ + This SubDomain grown by a number of points on each side of some of its + Dimensions, clipped to the Grid (see `DerivedSubDomain`). + + Parameters + ---------- + growth : dict of {Dimension: int} + Number of points to grow by, per root Dimension. + """ + return DerivedSubDomain(self, growth) + + def _indicator(self, dim, offset): + """ + 1 if the point `offset` points away from the current one along `dim` lies + in this SubDomain, 0 otherwise, computed without branching. Always 1 if + this SubDomain spans all of `dim`. + + Parameters + ---------- + dim : Dimension + The root Dimension along which to test. + offset : expr-like + The offset, e.g. 2, or `i0` for a stencil in unexpanded form. + """ + maybe_subdim = self.dimension_map.get(dim, dim) + if not maybe_subdim.is_Sub: + return sympy.S.One + point = dim + offset + return (Max(0, Min(1, point - maybe_subdim.symbolic_min + 1)) * + Max(0, Min(1, maybe_subdim.symbolic_max - point + 1))) + + def _restrict(self, expr, exclude=()): + """ + `expr` restricted to this SubDomain at the evaluation point: zero at the + points outside of it, along all root Dimensions but those in `exclude`. + """ + for d in self.dimensions: + if d.root not in exclude: + expr = expr * self._indicator(d.root, 0) + return expr + @cached_property def _arg_names(self): try: @@ -755,6 +794,58 @@ def __setstate__(self, state): self._distributor = SubDistributor(self) +class DerivedSubDomain(SubDomain): + + """ + A SubDomain derived from a parent SubDomain, akin to a DerivedDimension: + the parent grown by a number of points on each side of some of its + Dimensions, clipped to the Grid. + + This is the region written by a finite-difference stencil with `halo=0` + applied to a field living on the parent SubDomain. + + Parameters + ---------- + parent : SubDomain + The SubDomain to grow. + growth : dict of {Dimension: int} + Number of points to grow by on each side, per root Dimension. + """ + + def __init__(self, parent, growth): + self._parent = parent + self._growth = frozendict(growth) + self.name = f"{parent.name}_grown" + super().__init__(grid=parent.grid) + + @property + def parent(self): + """The SubDomain this one is derived from.""" + return self._parent + + @property + def growth(self): + """ + Number of points by which this SubDomain extends its parent on each side, + per root Dimension. + """ + return self._growth + + def define(self, dimensions): + regions = {} + for d, v in self.parent.define(dimensions).items(): + radius = self.growth.get(d, 0) + if isinstance(v, Dimension) or radius == 0: + regions[d] = v + elif v[0] == 'middle': + ltkn, rtkn = max(v[1] - radius, 0), max(v[2] - radius, 0) + regions[d] = d if ltkn == rtkn == 0 else ('middle', ltkn, rtkn) + else: + side, thickness = v + regions[d] = (side, min(thickness + radius, self.grid.shape[d])) + return regions + + class MultiSubDomain(AbstractSubDomain): """ diff --git a/devito/types/tensor.py b/devito/types/tensor.py index 08cb95fe63..ba9bf51a01 100644 --- a/devito/types/tensor.py +++ b/devito/types/tensor.py @@ -12,6 +12,7 @@ from sympy.core.sympify import converter as sympify_converter from devito.finite_differences import Differentiable +from devito.finite_differences.differentiable import merge_growth from devito.finite_differences.tools import make_shift_x0 from devito.types.basic import AbstractTensor from devito.types.dense import Function, TimeFunction @@ -179,7 +180,7 @@ def _eval_at(self, func, **kwargs): """ def entries(i, j, func): return getattr(self[i, j], '_eval_at', - lambda x: self[i, j])(func[i, j], **kwargs) + lambda x, **kw: self[i, j])(func[i, j], **kwargs) entry = lambda i, j: entries(i, j, func) return self._new(self.rows, self.cols, entry) @@ -221,6 +222,16 @@ def is_diagonal(self): return np.all([self[i, j] == 0 for j in range(self.cols) for i in range(self.rows) if i != j]) + @cached_property + def _has_zero_halo(self): + """True if a component has derivatives with `halo=0`.""" + return any(x._has_zero_halo for x in self.values() if x != 0) + + @cached_property + def _growth(self): + """The growth of the components (see `Differentiable._growth`).""" + return merge_growth(x._growth for x in self.values() if x != 0) + def _evaluate(self, **kwargs): def _do_evaluate(x): try: diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index 46fad9cdeb..269e5b6e60 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -15,6 +15,7 @@ ) from devito.symbolics import indexify, retrieve_indexed from devito.types.dimension import StencilDimension +from devito.types.grid import DerivedSubDomain from devito.warnings import DevitoWarning _PRECISION = 9 @@ -1669,3 +1670,407 @@ def test_deriv_mixed_subdomain_cpml_adjoint(): lhs = np.vdot(out.data, out_bar.data) rhs = np.vdot(q.data, q_bar.data) assert np.isclose(lhs, rhs, rtol=1e-12) + + +class TestHaloZero: + """ + Derivatives with `halo=0` treat their argument as zero outside the SubDomain + of their equation. With `.T`, `Inc(g_bar, out_bar.dx(halo=0).T, subdomain=S)` + is the adjoint of `Eq(out, g.dx, subdomain=S)`. + """ + + @staticmethod + def dot_test(grid, subdomain, deriv, so, eq_type=Inc, opt='advanced'): + """ + Relative error of the dot test between `Eq(out, g., subdomain=S)` + and `eq_type(g_bar, out_bar.(halo=0).T, subdomain=S)`, both built + with the optimization options `opt`. + """ + kwargs = {'grid': grid, 'space_order': so, 'dtype': np.float64} + g = Function(name='g', **kwargs) + out = Function(name='out', **kwargs) + out_bar = Function(name='out_bar', **kwargs) + g_bar = Function(name='g_bar', **kwargs) + + rng = np.random.default_rng(1) + g.data[:] = rng.normal(size=g.shape) + # Nonzero outside S too: the adjoint must ignore it there + out_bar.data[:] = rng.normal(size=out_bar.shape) + + Operator(Eq(out, getattr(g, deriv), subdomain=subdomain), opt=opt).apply() + adjoint = getattr(out_bar, deriv)(halo=0).T + Operator(eq_type(g_bar, adjoint, subdomain=subdomain), opt=opt).apply() + + lhs = np.vdot(out.data, out_bar.data) + rhs = np.vdot(g.data, g_bar.data) + return abs(lhs - rhs) / abs(lhs) + + @pytest.mark.parametrize('so', [2, 4, 8]) + @pytest.mark.parametrize('deriv', ['dx', 'dx2']) + @pytest.mark.parametrize('region', [('left', 6), ('right', 5), ('middle', 7, 9)]) + def test_adjoint(self, region, deriv, so): + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': region}, grid=grid) + + assert self.dot_test(grid, subdomain, deriv, so) < 1e-12 + + def test_adjoint_eq(self): + """An assignment covers the whole grown region, like an increment.""" + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) + + assert self.dot_test(grid, subdomain, 'dx', 4, eq_type=Eq) < 1e-12 + + @pytest.mark.parametrize('deriv', ['dx', 'dx2']) + def test_adjoint_unexpanded(self, deriv): + """The stencil stays a loop over its points: the mask is read within it.""" + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) + opt = ('advanced', {'expand': False}) + + assert self.dot_test(grid, subdomain, deriv, 8, opt=opt) < 1e-12 + + out_bar = Function(name='out_bar', grid=grid, space_order=8) + g_bar = Function(name='g_bar', grid=grid, space_order=8) + op = Operator(Inc(g_bar, out_bar.dx(halo=0).T, subdomain=subdomain), opt=opt) + assert 'for (int i0' in str(op.ccode) + + @pytest.mark.parametrize('so', [2, 4, 8]) + @pytest.mark.parametrize('stagger_in, stagger_out', [(True, False), (False, True)]) + def test_adjoint_staggered(self, stagger_in, stagger_out, so): + """ + Staggered first derivatives, as in elastic CPMLs: with `g` and `out` + staggered apart, the adjoint of `Eq(out, g.dx, subdomain=S)` is + `-out_bar.dx(halo=0)` evaluated at the location of `g`. + """ + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + x = grid.dimensions[0] + subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) + kwargs = {'grid': grid, 'space_order': so, 'dtype': np.float64} + g = Function(name='g', staggered=x if stagger_in else None, **kwargs) + g_bar = Function(name='g_bar', staggered=x if stagger_in else None, **kwargs) + out = Function(name='out', staggered=x if stagger_out else None, **kwargs) + out_bar = Function(name='out_bar', staggered=x if stagger_out else None, + **kwargs) + + rng = np.random.default_rng(5) + g.data[:] = rng.normal(size=g.shape) + out_bar.data[:] = rng.normal(size=out_bar.shape) + + Operator(Eq(out, g.dx, subdomain=subdomain)).apply() + Operator(Inc(g_bar, -out_bar.dx(halo=0), subdomain=subdomain)).apply() + + lhs = np.vdot(out.data, out_bar.data) + rhs = np.vdot(g.data, g_bar.data) + assert np.isclose(lhs, rhs, rtol=1e-12) + + def test_adjoint_unexpanded_3d(self): + """ + In unexpanded form, the indicators vary along the stencil: they must not + be hoisted as invariants of the outer Dimensions, as on GPUs. + """ + grid = Grid(shape=(12, 21, 12), extent=(1., 1., 1.), dtype=np.float64) + subdomain = Region('s', {'y': ('right', 6)}, grid=grid) + opt = ('advanced', {'expand': False}) + + assert self.dot_test(grid, subdomain, 'dy', 8, opt=opt) < 1e-12 + + @pytest.mark.parametrize('deriv', ['dx', 'dy']) + 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) + + assert self.dot_test(grid, subdomain, deriv, 4) < 1e-12 + + def test_forward(self): + """`g.dx(halo=0)` is the derivative of `g` extended by zero outside S.""" + so = 4 + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) + kwargs = {'grid': grid, 'space_order': so, 'dtype': np.float64} + g = Function(name='g', **kwargs) + g_zero = Function(name='g_zero', **kwargs) + out = Function(name='out', **kwargs) + ref = Function(name='ref', **kwargs) + + g.data[:] = np.random.default_rng(2).normal(size=g.shape) + g_zero.data[7:-9] = g.data[7:-9] + + Operator([Eq(out, g.dx(halo=0), subdomain=subdomain), + Eq(ref, g_zero.dx)]).apply() + + assert np.linalg.norm(ref.data) > 0 + assert np.allclose(out.data, ref.data, rtol=1e-12, atol=1e-12) + + def test_default_unchanged(self): + """ + Without `halo=0`, `.T` in a SubDomain equation is the full-grid transpose + evaluated on S only. + """ + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + g_bar = Function(name='g_bar', **kwargs) + ref = Function(name='ref', **kwargs) + out_bar = Function(name='out_bar', **kwargs) + out_bar.data[:] = np.random.default_rng(3).normal(size=out_bar.shape) + + Operator([Inc(g_bar, out_bar.dx.T, subdomain=subdomain), + Inc(ref, out_bar.dx.T)]).apply() + + assert np.allclose(g_bar.data[7:-9], ref.data[7:-9], rtol=1e-12) + assert np.all(g_bar.data[:7] == 0) and np.all(g_bar.data[-9:] == 0) + + def test_same_subdomain_two_grids(self): + """Equal SubDomains on different Grids get their own masks.""" + for shape in [(24,), (24,)]: + grid = Grid(shape=shape, extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': ('left', 6)}, grid=grid) + + assert self.dot_test(grid, subdomain, 'dx', 4) < 1e-12 + + def test_transpose_keeps_halo(self): + grid = Grid(shape=(8,)) + f = Function(name='f', grid=grid, space_order=4) + + assert f.dx(halo=0).halo == 0 + assert f.dx(halo=0).T.halo == 0 + assert f.dx.T.halo is None + assert f.dx(halo=0) != f.dx + + def test_no_subdomain(self): + """Without a SubDomain, halo=0 has no effect.""" + grid = Grid(shape=(16,), extent=(15.,), dtype=np.float64) + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + f = Function(name='f', **kwargs) + g = Function(name='g', **kwargs) + ref = Function(name='ref', **kwargs) + f.data[:] = np.random.default_rng(4).normal(size=f.shape) + + Operator([Eq(g, f.dx(halo=0).T), Eq(ref, f.dx.T)]).apply() + + assert np.all(g.data == ref.data) + + @pytest.mark.parametrize('region, grown', [ + (('left', 6), ('left', 8)), + (('right', 5), ('right', 7)), + (('middle', 7, 9), ('middle', 5, 7)), + (('middle', 1, 9), ('middle', 0, 7)), + (('left', 15), ('left', 16)), + ]) + def test_derived_subdomain(self, region, grown): + """ + A DerivedSubDomain spans its parent grown by `growth`, within the Grid. + """ + grid = Grid(shape=(16,)) + x = grid.dimensions[0] + subdomain = Region('s', {'x': region}, grid=grid) + + grown_subdomain = subdomain._grow({x: 2}) + + assert isinstance(grown_subdomain, DerivedSubDomain) + assert grown_subdomain.parent is subdomain + assert grown_subdomain.define(grid.dimensions) == {x: grown} + + def test_terms(self): + """ + In a sum, the terms without halo=0 derivatives stay on the SubDomain, + while the others extend past it. A factor of a halo=0 derivative is not + restricted. + """ + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + f, g, c, out, ref = (Function(name=n, **kwargs) + for n in ('f', 'g', 'c', 'out', 'ref')) + rng = np.random.default_rng(8) + for h in (f, g, c): + h.data[:] = rng.normal(size=h.shape) + + Operator(Eq(out, f + c*g.dx(halo=0), subdomain=subdomain)).apply() + Operator([Eq(ref, f, subdomain=subdomain), + Inc(ref, c*g.dx(halo=0), subdomain=subdomain)]).apply() + + assert np.linalg.norm(ref.data[:7]) > 0 + assert np.allclose(out.data, ref.data, rtol=1e-12, atol=1e-12) + + def test_sum_factor(self): + """ + A sum without halo=0 derivatives, as a factor of one, is not restricted: + the product extends past the SubDomain. + """ + grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) + subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + f, g, c, d, out = (Function(name=n, **kwargs) + for n in ('f', 'g', 'c', 'd', 'out')) + rng = np.random.default_rng(9) + for h in (f, g, c): + h.data[:] = rng.normal(size=h.shape) + + Operator(Eq(d, g.dx(halo=0), subdomain=subdomain)).apply() + Operator(Eq(out, (f + c)*g.dx(halo=0), subdomain=subdomain)).apply() + + expected = (f.data + c.data)*d.data + assert np.linalg.norm(expected[:7]) > 0 + assert np.allclose(out.data, expected, rtol=1e-12, atol=1e-12) + + def test_vector(self): + """ + In a vector equation, each component is restricted and extended past the + SubDomain as in the corresponding scalar equation. + """ + grid = Grid(shape=(12, 14), extent=(11., 13.), dtype=np.float64) + subdomain = Region('s', {'x': ('left', 5)}, grid=grid) + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + f = Function(name='f', **kwargs) + c = VectorFunction(name='c', staggered=(None, None), **kwargs) + w = VectorFunction(name='w', staggered=(None, None), **kwargs) + ref = [Function(name=f'ref{i}', **kwargs) for i in range(2)] + rng = np.random.default_rng(10) + for h in (f, *c): + h.data[:] = rng.normal(size=h.shape) + + Operator(Eq(w, c*f.dx(halo=0), subdomain=subdomain)).apply() + Operator([Eq(r, ci*f.dx(halo=0), subdomain=subdomain) + for r, ci in zip(ref, c, strict=True)]).apply() + + for wi, r in zip(w, ref, strict=True): + assert np.linalg.norm(r.data[5:]) > 0 + assert np.allclose(wi.data, r.data, rtol=1e-12, atol=1e-12) + + def test_vector_staggered(self): + """ + Dot test of a staggered vector equation restricted to a SubDomain, as in + elastic CPMLs: the adjoint of `Eq(v, grad(p), subdomain=S)`, with `v` + staggered, is `Inc(p_bar, -div(v_bar))` with halo=0 derivatives. + """ + grid = Grid(shape=(14, 16), extent=(13., 15.), dtype=np.float64) + x, y = grid.dimensions + subdomain = Region('s', {'x': ('left', 5), 'y': ('middle', 4, 5)}, grid=grid) + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + p = Function(name='p', **kwargs) + p_bar = Function(name='p_bar', **kwargs) + v = VectorFunction(name='v', **kwargs) + v_bar = VectorFunction(name='v_bar', **kwargs) + rng = np.random.default_rng(11) + p.data[:] = rng.normal(size=p.shape) + for vi in v_bar: + vi.data[:] = rng.normal(size=vi.shape) + + Operator(Eq(v, grad(p), subdomain=subdomain)).apply() + Operator(Inc(p_bar, -(v_bar[0].dx(halo=0) + v_bar[1].dy(halo=0)), + subdomain=subdomain)).apply() + + lhs = sum(np.vdot(vi.data, vi_bar.data) + for vi, vi_bar in zip(v, v_bar, strict=True)) + rhs = np.vdot(p.data, p_bar.data) + assert np.isclose(lhs, rhs, rtol=1e-12) + + def test_unrestricted_dimension(self): + """halo=0 along a Dimension the SubDomain spans has no effect.""" + grid = Grid(shape=(12, 14), extent=(11., 13.), dtype=np.float64) + subdomain = Region('s', {'x': ('left', 5)}, grid=grid) + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + f = Function(name='f', **kwargs) + g = Function(name='g', **kwargs) + ref = Function(name='ref', **kwargs) + f.data[:] = np.random.default_rng(6).normal(size=f.shape) + + Operator([Eq(g, f.dy(halo=0), subdomain=subdomain), + Eq(ref, f.dy, subdomain=subdomain)]).apply() + + assert np.all(g.data == ref.data) + + def test_invalid_halo(self): + grid = Grid(shape=(8,)) + f = Function(name='f', grid=grid, space_order=4) + + with pytest.raises(ValueError): + f.dx(halo=1) + + def test_nested_halo(self): + """Differentiating a derivative with `halo=0` is not supported.""" + grid = Grid(shape=(24,)) + subdomain = Region('s', {'x': ('left', 6)}, grid=grid) + f = Function(name='f', grid=grid, space_order=4) + g = Function(name='g', grid=grid, space_order=4) + + with pytest.raises(NotImplementedError): + _ = Eq(g, f.dx(halo=0).dx(halo=0), subdomain=subdomain).evaluate + + def test_growth(self): + """ + The growth of an evaluated expression is the largest stencil radius of + its `halo=0` derivatives, per root Dimension. + """ + grid = Grid(shape=(24, 24)) + x, y = grid.dimensions + subdomain = Region('s', {'x': ('left', 6), 'y': ('left', 6)}, grid=grid) + f = Function(name='f', grid=grid, space_order=8) + g = Function(name='g', grid=grid, space_order=8) + + rhs = (f.dx(halo=0) + f.dy(fd_order=4, halo=0) + f.dx(fd_order=2)) + eq = Eq(g, rhs, subdomain=subdomain).evaluate + + assert rhs._has_zero_halo + assert eq.subdomain.growth == {x: 4, y: 2} + + def test_cpml_adjoint(self): + """ + Dot test of a one-face CPML derivative written on the CPML strip S only. + + forward: out = D q + psi[n+1] on S + psi[n+1] = b psi[n] + a D q on S + adjoint: q_bar += D^T (out_bar + a (psi_bar + out_bar)), + with halo=0 and subdomain=S + + The adjoint needs neither a zero-padded work field nor a mask. + """ + nx, width, so, nt = 32, 10, 8, 8 + grid = Grid(shape=(nx,), extent=(float(nx - 1),), dtype=np.float64) + time = grid.time_dim + x = grid.dimensions[0] + strip = Region('strip', {'x': ('left', width)}, grid=grid) + + kwargs = {'grid': grid, 'save': nt, 'space_order': so, 'dtype': np.float64} + q = TimeFunction(name='q', **kwargs) + out = TimeFunction(name='out', **kwargs) + q_bar = TimeFunction(name='q_bar', **kwargs) + out_bar = TimeFunction(name='out_bar', **kwargs) + psi = TimeFunction(name='psi', grid=strip, time_order=1, space_order=so, + dtype=np.float64) + # The adjoint stencil reads psi_bar one radius past the grown strip + psi_bar = TimeFunction(name='psi_bar', grid=strip, time_order=1, + space_order=(so, so, so), dtype=np.float64) + a = Function(name='a', grid=grid, dimensions=(x,), shape=(nx,), + space_order=so, dtype=np.float64) + b = Function(name='b', grid=grid, dimensions=(x,), shape=(nx,), + dtype=np.float64) + a.data[:6] = -0.1 + b.data[:] = 1. + b.data[:6] = 0.8 + + forward = Operator([ + Eq(psi.forward, b * psi + a * q.dx, subdomain=strip), + Eq(out, q.dx + psi.forward, subdomain=strip) + ]) + value = psi_bar + out_bar + adjoint = Operator([ + Eq(psi_bar.backward, b * value, subdomain=strip, implicit_dims=(time,)), + Inc(q_bar, (out_bar + a * value).dx(halo=0).T, subdomain=strip, + implicit_dims=(time,)) + ]) + + # A single time loop + assert str(adjoint.ccode).count('for (int time') == 1 + + rng = np.random.default_rng(7) + q.data[:] = rng.normal(size=q.shape) + out_bar.data[:] = rng.normal(size=out_bar.shape) + forward.apply(time_m=0, time_M=nt - 1) + adjoint.apply(time_m=0, time_M=nt - 1) + + lhs = np.vdot(out.data, out_bar.data) + rhs = np.vdot(q.data, q_bar.data) + assert np.isclose(lhs, rhs, rtol=1e-12)