api: introduce nearset neighboor interpolation and assignment injection for it

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,22 +443,18 @@ 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
         last = [ecls(lhs, rhs, implicit_dims=implicit_dims)]
 
         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

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,14 +1161,17 @@ 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:
             field = field.subs({field.time_dim: u_t})
         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):

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
 # ---------------------------------------------------------------------------