Skip to content

Commit f29efcc

Browse files
committed
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. - Each stencil tap is multiplied by a SubDomainIndicator (SubDomain.indicator) of the grid point it reads: 1 inside the SubDomain, 0 outside. It records the stencil radius (IndexSet.radius) and works for expanded, unexpanded and staggered derivatives. - Eq._extend evaluates such equations over the SubDomain grown by the largest radius (SubDomain.grow, a GrownSubDomain), restricts the other terms to the SubDomain at the evaluation point, and expands the indicators into branch-free integer MIN/MAX, which invariant hoisting computes once. - 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; staggered dot tests against `-out_bar.dx(halo=0)`, as in elastic adjoints; a forward check; and a one-face CPML operator.
1 parent 66418f7 commit f29efcc

6 files changed

Lines changed: 550 additions & 16 deletions

File tree

‎devito/finite_differences/derivative.py‎

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ def _fd_priority(self):
102102

103103
__rargs__ = ('expr', '*dims')
104104
__rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs',
105-
'x0', 'method', 'weights')
105+
'x0', 'method', 'weights', 'halo')
106106

107107
def __new__(cls, expr, *dims, **kwargs):
108108
# Validate the input arguments `expr`, `dims` and `deriv_order`
@@ -161,6 +161,7 @@ def __new__(cls, expr, *dims, **kwargs):
161161
obj._transpose = kwargs.get("transpose", direct)
162162
obj._method = kwargs.get("method", 'FD')
163163
obj._weights = cls._process_weights(**kwargs)
164+
obj._halo = cls._validate_halo(kwargs.get("halo"))
164165

165166
ppsubs = kwargs.get("subs", kwargs.get("_ppsubs", []))
166167
processed = []
@@ -177,6 +178,16 @@ def __new__(cls, expr, *dims, **kwargs):
177178

178179
return obj
179180

181+
@staticmethod
182+
def _validate_halo(halo):
183+
"""
184+
Validate `halo`. Only None (read the argument everywhere) and 0 (treat
185+
the argument as zero outside the equation's SubDomain) are supported.
186+
"""
187+
if halo not in (None, 0):
188+
raise ValueError(f"Expected halo=None or halo=0, got halo={halo}")
189+
return halo
190+
180191
@staticmethod
181192
def _validate_expr(expr):
182193
"""
@@ -325,6 +336,8 @@ def _process_weights(cls, **kwargs):
325336
def __call__(self, x0=None, fd_order=None, side=None, method=None, **kwargs):
326337
weights = kwargs.get('weights', kwargs.get('w'))
327338
rkw = {}
339+
if 'halo' in kwargs:
340+
rkw['halo'] = kwargs['halo']
328341
if side is not None:
329342
rkw['side'] = side
330343
if method is not None:
@@ -457,6 +470,14 @@ def side(self):
457470
def transpose(self):
458471
return self._transpose
459472

473+
@property
474+
def halo(self):
475+
"""
476+
None if the argument is read everywhere, 0 if it is treated as zero
477+
outside the SubDomain of the equation the Derivative belongs to.
478+
"""
479+
return self._halo
480+
460481
@property
461482
def is_TimeDependent(self):
462483
return self.expr.is_TimeDependent
@@ -581,6 +602,13 @@ def _eval_fd(self, expr, **kwargs):
581602
if expr.is_Add and any(len(indices_at(expr, d)) > 1 for d in self.dims):
582603
return expr.func(*[self._eval_fd(a, **kwargs) for a in expr.args])
583604

605+
# The SubDomain mask read by a halo=0 derivative can't be shifted again
606+
if any(d.halo is not None for d in expr.find(Derivative)):
607+
raise NotImplementedError(
608+
f"{self} differentiates a derivative with halo=0, which is not "
609+
"supported"
610+
)
611+
584612
# Step 1: Evaluate non-derivative x0. We currently enforce a simple 2nd order
585613
# interpolation to avoid very expensive finite differences on top of it
586614
x0_deriv = self._filter_dims(self.x0)
@@ -598,6 +626,15 @@ def _eval_fd(self, expr, **kwargs):
598626
# otherwise an IndexSum will returned
599627
expand = kwargs.get('expand', True)
600628

629+
# With halo=0, `expr` is treated as zero outside the equation's SubDomain
630+
subdomain = kwargs.get('subdomain') if self.halo is not None else None
631+
if subdomain is not None and (subdomain.is_MultiSubDomain or
632+
self.method != 'FD'):
633+
raise NotImplementedError(
634+
f"halo=0 is only supported with method='FD' on a SubDomain, not "
635+
f"with method={self.method} on {subdomain}"
636+
)
637+
601638
# Step 3: Evaluate FD of the new expression
602639
if self.method == 'RSFD':
603640
assert len(self.dims) == 1
@@ -607,13 +644,15 @@ def _eval_fd(self, expr, **kwargs):
607644
assert self.method == 'FD'
608645
res = cross_derivative(expr, self.dims, self.fd_order, self.deriv_order,
609646
matvec=self.transpose, x0=x0_deriv, expand=expand,
610-
side=self.side, weights=self.weights)
647+
side=self.side, weights=self.weights,
648+
subdomain=subdomain)
611649
else:
612650
assert self.method == 'FD'
613651
res = generic_derivative(expr, self.dims[0], self.fd_order[0],
614652
self.deriv_order[0], weights=self.weights,
615653
side=self.side, matvec=self.transpose,
616-
x0=self.x0, expand=expand)
654+
x0=self.x0, expand=expand,
655+
subdomain=subdomain)
617656

618657
# Step 4: Apply substitutions
619658
for e in self._ppsubs:

‎devito/finite_differences/finite_difference.py‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,8 @@ def index_at(expr, dim):
125125

126126
@check_input
127127
def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
128-
coefficients='taylor', expand=True, weights=None, side=None):
128+
coefficients='taylor', expand=True, weights=None, side=None,
129+
subdomain=None):
129130
"""
130131
Arbitrary-order derivative of a given expression.
131132
@@ -151,6 +152,10 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
151152
expand : bool, optional, default=True
152153
If True, the derivative is fully expanded as a sum of products,
153154
otherwise an IndexSum is returned.
155+
subdomain : SubDomain, optional
156+
If given, `expr` is treated as zero outside `subdomain` (see
157+
`Derivative`'s `halo`): each stencil tap is multiplied by the
158+
indicator of `subdomain` at the point it reads.
154159
155160
Returns
156161
-------
@@ -171,7 +176,7 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
171176
coefficients = 'taylor' if dim.is_Time else expr.coefficients
172177

173178
return make_derivative(expr, dim, fd_order, deriv_order, side,
174-
matvec, x0, coefficients, expand, weights)
179+
matvec, x0, coefficients, expand, weights, subdomain)
175180

176181

177182
# Backward compatibility
@@ -180,7 +185,7 @@ def first_derivative(expr, dim, fd_order, **kwargs):
180185

181186

182187
def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coefficients,
183-
expand, weights=None):
188+
expand, weights=None, subdomain=None):
184189
# Always expand time derivatives to avoid issue with buffering and streaming.
185190
# Time derivative are almost always short stencils and won't benefit from
186191
# unexpansion in the rare case the derivative is not evaluated for time stepping.
@@ -222,13 +227,20 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
222227
if callable(expand):
223228
expand = expand(dim)
224229

230+
# With a `subdomain`, `expr` is treated as zero outside it: every stencil tap
231+
# is multiplied by the indicator of the point it reads, 1 if `subdomain`
232+
# spans the whole of `dim`
233+
225234
if not expand and indices.expr is not None:
226235
weights = Weights(name='w', dimensions=indices.free_dim,
227236
initvalue=weights, dtype=expr.dtype)
228237

229238
# Inject the StencilDimension
230239
# E.g. `x + i*h_x` into `f(x)` s.t. `f(x + i*h_x)`
231240
expr = expr.shift(dim, indices.expr - dim)
241+
if subdomain is not None:
242+
offset = indices.offset(indices.expr)
243+
expr = expr * subdomain.indicator(dim.root, offset, indices.radius)
232244

233245
# Re-evaluate any off-the-grid Functions potentially impacted by the FD
234246
# unless a pure number
@@ -243,6 +255,9 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
243255
for i, c in zip(indices, weights, strict=True):
244256
# The FD term
245257
term = expr.shift(dim, i - dim) * c
258+
if subdomain is not None:
259+
term = term * subdomain.indicator(dim.root, indices.offset(i),
260+
indices.radius)
246261

247262
# Re-evaluate any off-the-grid Functions potentially impacted by the FD
248263
# unless a pure number

‎devito/finite_differences/tools.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from functools import partial, wraps
1+
from functools import cached_property, partial, wraps
22
from itertools import product
33

44
import numpy as np
@@ -177,6 +177,18 @@ def __repr__(self):
177177
def spacing(self):
178178
return self.dim.spacing
179179

180+
def offset(self, index):
181+
"""
182+
Offset, in grid points, of the point `index` from `dim`: e.g. 2 for
183+
`x + 2*h_x`, or `i0` for `x + i0*h_x` in unexpanded form.
184+
"""
185+
return sympify((index - self.dim) / self.spacing)
186+
187+
@cached_property
188+
def radius(self):
189+
"""Largest distance, in grid points, between `dim` and the points."""
190+
return max(abs(self.offset(i)) for i in self)
191+
180192
def transpose(self):
181193
"""
182194
Transpose the IndexSet.

‎devito/types/equation.py‎

Lines changed: 62 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,18 +125,79 @@ def _evaluate(self, **kwargs):
125125
"""
126126
if self._interp_mode is not None:
127127
kwargs['interp_mode'] = self._interp_mode
128+
129+
subdomain = self.subdomain
130+
restrict = subdomain is not None and self._has_halo
131+
if restrict:
132+
# Derivatives with halo=0 treat their argument as zero outside it
133+
kwargs = {**kwargs, 'subdomain': subdomain}
134+
128135
try:
129136
lhs = self.lhs._evaluate(**kwargs)
130137
rhs = self.rhs._eval_at(self.lhs, **kwargs)._evaluate(**kwargs)
131138
except AttributeError:
132139
lhs, rhs = self._evaluate_args(**kwargs)
133-
eq = self.func(lhs, rhs, subdomain=self.subdomain,
140+
141+
if restrict:
142+
# Derivatives with halo=0 extend the rhs past the SubDomain
143+
rhs, subdomain = self._extend(rhs, subdomain)
144+
145+
eq = self.func(lhs, rhs, subdomain=subdomain,
134146
coefficients=self.substitutions,
135147
implicit_dims=self._implicit_dims,
136148
interp_mode=self._interp_mode)
137149

138150
return eq
139151

152+
@cached_property
153+
def _has_halo(self):
154+
"""True if the rhs has derivatives with halo=0."""
155+
from devito.symbolics import retrieve_derivatives # noqa
156+
157+
derivatives = retrieve_derivatives(self.rhs, mode='unique', deep=True)
158+
return any(d.halo is not None for d in derivatives)
159+
160+
@staticmethod
161+
def _extend(rhs, subdomain):
162+
"""
163+
The evaluated `rhs` restricted to `subdomain`, and the region over which
164+
it is evaluated.
165+
166+
The derivatives with halo=0 in `rhs` multiply their taps by the
167+
SubDomainIndicators of `subdomain`, and are nonzero up to their stencil
168+
radius past it. The region is `subdomain` grown by the largest radius.
169+
"""
170+
from devito.symbolics import search # noqa
171+
172+
# The indicators of each term; a term without indicators has no halo=0
173+
# derivative
174+
terms = sympy.Add.make_args(rhs)
175+
found = [search(t, lambda e: getattr(e, 'is_SubDomainIndicator', False))
176+
for t in terms]
177+
indicators = set().union(*found)
178+
179+
# None if the halo=0 derivatives are along Dimensions `subdomain` spans
180+
if not indicators:
181+
return rhs, subdomain
182+
183+
# The region: the SubDomain grown by the largest stencil radius
184+
radius = {}
185+
for i in indicators:
186+
radius[i.dim] = max(radius.get(i.dim, 0), i.radius)
187+
region = subdomain.grow(radius)
188+
189+
# The terms without indicators are restricted to the SubDomain at the
190+
# evaluation point
191+
here = [subdomain.indicator(d, 0, r) for d, r in radius.items()]
192+
terms = [t if f else sympy.Mul(*here, t)
193+
for t, f in zip(terms, found, strict=True)]
194+
rhs = rhs.func(*terms) if rhs.is_Add else terms[0]
195+
196+
# Finally, all indicators are expanded into branch-free expressions
197+
rhs = rhs.xreplace({i: i.branchless for i in indicators.union(here)})
198+
199+
return rhs, region
200+
140201
@property
141202
def _flatten(self):
142203
"""

0 commit comments

Comments
 (0)