diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index 1bba546635..8bc8535e59 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -21,7 +21,8 @@ from devito.types import CustomDimension, Eq, Evaluable, Inc, SubFunction, Symbol from devito.types.utils import DimensionTuple -__all__ = ['LinearInterpolator', 'PrecomputedInterpolator', 'SincInterpolator'] +__all__ = ['LinearInterpolator', 'NearestInterpolator', + 'PrecomputedInterpolator', 'SincInterpolator'] def check_radius(func): @@ -163,20 +164,23 @@ class Injection(UnevaluatedSparseOperation): Evaluates to a list of Eq objects. """ - __rargs__ = ('field', 'expr', 'implicit_dims') + UnevaluatedSparseOperation.__rargs__ + __rargs__ = ('field', 'expr', 'increment', 'implicit_dims') + \ + UnevaluatedSparseOperation.__rargs__ - def __new__(cls, field, expr, implicit_dims, interpolator): + def __new__(cls, field, expr, increment, implicit_dims, interpolator): obj = super().__new__(cls, interpolator) # TODO: unused now, but will be necessary to compute the adjoint obj.field = field obj.expr = expr + obj.increment = increment obj.implicit_dims = implicit_dims return obj def operation(self, **kwargs): return self.interpolator._inject(expr=self.expr, field=self.field, + increment=self.increment, implicit_dims=self.implicit_dims) def __repr__(self): @@ -353,6 +357,20 @@ def _interp_idx(self, variables, implicit_dims=None, subdomain=None, return idx_subs, temps + def _local_accumulator(self, expr, idx_subs, implicit_dims=None, subdomain=None): + """ + Generate a local accumulator for the interpolation/injection operation. + """ + # Accumulate point-wise contributions into a temporary + rhs = Symbol(name=f'sum{self.sfunction.name}', dtype=self.sfunction.dtype) + summands = [Eq(rhs, 0., implicit_dims=implicit_dims)] + # Substitute coordinate base symbols into the interpolation coefficients + weights = self._weights(subdomain=subdomain) + summands.extend([Inc(rhs, (weights * expr).xreplace(idx_subs), + implicit_dims=implicit_dims)]) + + return summands, rhs + @check_radius @check_coords def interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None): @@ -376,7 +394,7 @@ def interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None) @check_radius @check_coords - def inject(self, field, expr, implicit_dims=None): + def inject(self, field, expr, increment=True, implicit_dims=None): """ Generate equations injecting an arbitrary expression into a field. @@ -391,7 +409,7 @@ def inject(self, field, expr, implicit_dims=None): injection expression, but that should be honored when constructing the operator. """ - return Injection(field, expr, implicit_dims, self) + return Injection(field, expr, increment, implicit_dims, self) def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None): """ @@ -425,14 +443,10 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None idx_subs, temps = self._interp_idx(variables, implicit_dims=implicit_dims, subdomain=subdomain) - # Accumulate point-wise contributions into a temporary - rhs = Symbol(name=f'sum{self.sfunction.name}', dtype=self.sfunction.dtype) - summands = [Eq(rhs, 0., implicit_dims=implicit_dims)] - # Substitute coordinate base symbols into the interpolation coefficients - weights = self._weights(subdomain=subdomain) - summands.extend([Inc(rhs, (weights * expr).xreplace(idx_subs), - implicit_dims=implicit_dims)]) - + # Local scalar for accumulation over radius + summands, rhs = self._local_accumulator(expr, idx_subs, + implicit_dims=implicit_dims, + subdomain=subdomain) # Write/Incr `self` lhs = self.sfunction.subs(self_subs) ecls = Inc if increment else Eq @@ -440,7 +454,7 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None return temps + summands + last - def _inject(self, field, expr, implicit_dims=None): + def _inject(self, field, expr, increment=True, implicit_dims=None): """ Generate equations injecting an arbitrary expression into a field. @@ -494,9 +508,10 @@ def _inject(self, field, expr, implicit_dims=None): # Move all temporaries inside inner loop to improve parallelism # Can only be done for inject as interpolation needs a summing temp # that wouldn't allow collapsing - implicit_dims = implicit_dims + tuple(r.parent for r in - self._rdim(subdomain=subdomain, - shifts=shifts)) + with suppress(AttributeError): + implicit_dims = implicit_dims + tuple(r.parent for r in + self._rdim(subdomain=subdomain, + shifts=shifts)) # List of indirection indices for all adjacent grid points idx_subs, _temps = self._interp_idx(list(g_fields) + variables, @@ -505,8 +520,9 @@ def _inject(self, field, expr, implicit_dims=None): w = self._weights(subdomain=subdomain, shifts=shifts) temps.extend(_temps) - eqns.extend([Inc(f.xreplace(idx_subs), (w * e).xreplace(idx_subs), - implicit_dims=implicit_dims) + ecls = Inc if increment else Eq + eqns.extend([ecls(f.xreplace(idx_subs), (w * e).xreplace(idx_subs), + implicit_dims=implicit_dims) for f, e in zip(g_fields, g_exprs, strict=True)]) return filter_ordered(temps) + eqns @@ -713,6 +729,36 @@ def _arg_defaults(self, coords=None, sfunc=None, origin=None): return args +class NearestInterpolator(LinearInterpolator): + """ + Nearest neighbor interpolation scheme. + + Gridpoints and per-dim `(1-frac, frac)` weights are precomputed on the + host in fp64 (see `_arg_defaults`) and passed to the kernel as int32/fp + SubFunctions. The generated C only indexes those tables and never sees + `(c-o)/h` or `floor` on fp32. + """ + + _name = 'nearest' + + def _local_accumulator(self, expr, idx_subs, implicit_dims=None, subdomain=None): + return [], expr.xreplace(idx_subs) + + @memoized_meth + def _rdim(self, subdomain=None, shifts=None): + return DimensionTuple(*[0 for _ in self._cdim], getters=self._gdims) + + @memoized_meth + def _weights(self, subdomain=None, shifts=None): + return sympy.S.One + + def _gridpoints(self, shifts=None): + return self._generate_coeffs(tuple(shifts) if shifts else None)[0] + + def _coeffs(self, shifts=None): + return [] + + class PrecomputedInterpolator(WeightedInterpolator): """ Concrete implementation of WeightedInterpolator implementing a Precomputed diff --git a/devito/types/sparse.py b/devito/types/sparse.py index ccb330b2d6..d5952ef750 100644 --- a/devito/types/sparse.py +++ b/devito/types/sparse.py @@ -9,7 +9,7 @@ from devito.finite_differences import generate_fd_shortcuts from devito.mpi import MPI, SparseDistributor from devito.operations import ( - LinearInterpolator, PrecomputedInterpolator, SincInterpolator + LinearInterpolator, NearestInterpolator, PrecomputedInterpolator, SincInterpolator ) from devito.symbolics import indexify, retrieve_function_carriers from devito.tools import ( @@ -34,8 +34,9 @@ ] -_interpolators = {'linear': LinearInterpolator, 'sinc': SincInterpolator} -_default_radius = {'linear': 1, 'sinc': 4} +_interpolators = {'linear': LinearInterpolator, 'sinc': SincInterpolator, + 'nearest': NearestInterpolator} +_default_radius = {'linear': 1, 'sinc': 4, 'nearest': 0} class SparseSubFunction(SubFunction): @@ -380,6 +381,8 @@ def _pos_symbols(self, shifts=None): @cached_property def _point_increments(self): """Index increments in each Dimension for each point symbol.""" + if self.r == 0: + return ((0,) * self.grid.dim,) return tuple(product(range(-self.r+1, self.r+1), repeat=self.grid.dim)) @cached_property @@ -439,11 +442,14 @@ def interpolate(self, *args, **kwargs): """ return self.interpolator.interpolate(*args, **kwargs) - def inject(self, *args, **kwargs): + def inject(self, *args, increment=True, **kwargs): """ Implement an injection operation from a sparse point onto the grid """ - return self.interpolator.inject(*args, **kwargs) + if not increment and self.interpolation != 'nearest': + raise ValueError("Assignment injection is only supported" + "for nearest-neighbor interpolation") + return self.interpolator.inject(*args, increment=increment, **kwargs) def guard(self, expr=None): """ @@ -982,6 +988,8 @@ def __interp_setup__(self, interpolation='linear', r=None, **kwargs): raise ValueError("'sinc' interpolator requires a radius of at most 10") elif interpolation == 'linear' and self._radius != 1: self._radius = 1 + elif interpolation == 'nearest': + self._radius = 0 @cached_property def _coordinate_symbols(self): @@ -1135,7 +1143,7 @@ def interpolate(self, expr, u_t=None, p_t=None, increment=False, implicit_dims=N return super().interpolate(expr, increment=increment, self_subs=subs, implicit_dims=implicit_dims) - def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None): + def inject(self, field, expr, increment=True, u_t=None, p_t=None, implicit_dims=None): """ Generate equations injecting an arbitrary expression into a field. @@ -1153,6 +1161,8 @@ def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None): An ordered list of Dimensions that do not explicitly appear in the injection expression, but that should be honored when constructing the operator. + increment: bool, optional + If Flalse, generate assignment (Eq) rather than increment (Inc). """ # Apply optional time symbol substitutions to field and expr if u_t is not None: @@ -1160,7 +1170,8 @@ def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None): if p_t is not None: expr = expr.subs({self.time_dim: p_t}) - return super().inject(field, expr, implicit_dims=implicit_dims) + return super().inject(field, expr, increment=increment, + implicit_dims=implicit_dims) @property def forward(self): diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 5d8142a412..fdd69fc1d7 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -11,7 +11,9 @@ SparseFunction, SparseTimeFunction, SubDomain, TimeFunction, VectorFunction, switchconfig ) -from devito.operations.interpolators import LinearInterpolator, SincInterpolator +from devito.operations.interpolators import ( + LinearInterpolator, SincInterpolator, _cell_indices +) from devito.tools import as_tuple from examples.seismic import ( AcquisitionGeometry, Receiver, RickerSource, TimeAxis, demo_model @@ -48,11 +50,12 @@ def unit_box_time(name='a', shape=(11, 11), space_order=1): return a -def points(grid, ranges, npoints, name='points'): +def points(grid, ranges, npoints, name='points', interpolation='linear'): """Create a set of sparse points from a set of coordinate ranges for each spatial dimension. """ - points = SparseFunction(name=name, grid=grid, npoint=npoints) + points = SparseFunction(name=name, grid=grid, npoint=npoints, + interpolation=interpolation) for i, r in enumerate(ranges): points.coordinates.data[:, i] = np.linspace(r[0], r[1], npoints) return points @@ -524,6 +527,89 @@ def test_inject_staggered_mixed(self): # Should be a single loop nest with 3 injections assert_structure(op, ['p_p,rp_px,rp_py,rp_pz'], 'p_prp_pxrp_py,rp_pz') + def test_assign_not_supported(self): + grid = Grid((11, 11, 11)) + a = Function(name='a', grid=grid, space_order=2) + p = SparseFunction(name="p", grid=grid, nt=10, npoint=1) + + with pytest.raises(ValueError): + p.inject(a, expr=1., increment=False) + + +# --------------------------------------------------------------------------- +# Nearest interpolation / injection +# --------------------------------------------------------------------------- + + +class TestNearest: + + """Tests for NearestSparseFunction / NearestSparseTimeFunction.""" + + @pytest.mark.parametrize('shape, coords', SHAPE_COORDS) + def test_nearest_interpolation(self, shape, coords, npoints=20): + """Test interpolation with NearestSparseFunction which uses nearest + neighbour interpolation. + """ + a = unit_box(shape=shape) + p = points(a.grid, ranges=coords, npoints=npoints, interpolation='nearest') + # nearest grdpoint + xcoords = _cell_indices(p.coordinates.data, a.grid, None, + a.grid.spacing, a.grid.origin)[:, 0] * a.grid.spacing[0] + + expr = p.interpolate(a) + op = Operator(expr) + + op(a=a) + assert np.allclose(p.data[:], xcoords, rtol=1e-6) + + @pytest.mark.parametrize('shape, coords, result', SHAPE_COORDS_INJECT) + def test_nearest_injection(self, shape, coords, result, npoints=19): + """Test injection with NearestSparseFunction which uses nearest + neighbour injection. + """ + a = unit_box(shape=shape) + a.data[:] = 0. + p = points(a.grid, ranges=coords, npoints=npoints, interpolation='nearest') + + expr = p.inject(a, Float(1.)) + + op = Operator(expr) + + op(a=a) + + d64 = np.array([.1]*a.grid.dim, dtype=np.float64) + o64 = np.array([0]*a.grid.dim, dtype=np.float64) + indices = _cell_indices(p.coordinates.data, a.grid, None, + d64, o64) + pos, pcount = np.unique(indices, return_counts=True, axis=0) + + assert np.sum(a.data) == npoints + for p, pc in zip(pos, pcount, strict=True): + assert a.data[tuple(p)] == pc + + @pytest.mark.parametrize('shape, coords, result', SHAPE_COORDS_INJECT) + def test_nearest_injection_assign(self, shape, coords, result, npoints=19): + a = unit_box(shape=shape) + a.data[:] = 0. + p = points(a.grid, ranges=coords, npoints=npoints, interpolation='nearest') + + expr = p.inject(a, Float(1.), increment=False) + + op = Operator(expr) + + op(a=a) + + d64 = np.array([.1]*a.grid.dim, dtype=np.float64) + o64 = np.array([0]*a.grid.dim, dtype=np.float64) + indices = _cell_indices(p.coordinates.data, a.grid, None, + d64, o64) + pos, _ = np.unique(indices, return_counts=True, axis=0) + + assert np.sum(a.data) == len(pos) + for i in indices: + assert a.data[tuple(i)] == 1 + + # --------------------------------------------------------------------------- # Precomputed interpolation / injection # ---------------------------------------------------------------------------