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 1e395e1a38..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 @@ -329,8 +349,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): @@ -530,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 @@ -630,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 @@ -1318,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 @@ -1331,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 @@ -1351,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 @@ -1363,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 73becb8b84..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,13 +227,21 @@ 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) # 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) + 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._subs(dim, i) * c + 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 22f47b5c94..4574321de3 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,19 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None): terms = cbk_compose(i) + # 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: + 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 +1625,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 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 b5169efbcd..269e5b6e60 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 ( @@ -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 @@ -1566,3 +1567,510 @@ 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) + + +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)