api: fix injection witjh staggering

Author: mloubout

devito/operations/interpolators.py

@@ -230,7 +230,7 @@ def r(self):
         return self.sfunction.r
 
     @memoized_meth
-    def _weights(self, subdomain=None):
+    def _weights(self, subdomain=None, shifts=None):
         raise NotImplementedError
 
     @property
@@ -243,8 +243,24 @@ def _cdim(self):
         dims = [self.sfunction._crdim(d) for d in self._gdims]
         return dims
 
+    def _field_shifts(self, field):
+        """
+        Per-grid-Dimension half-cell shift induced by ``field``'s staggering
+        (e.g. ``h_x/2`` for a field staggered in ``x``). Returns None for
+        unstaggered fields. SubDomain-induced origin offsets are deliberately
+        ignored — they are not staggering.
+        """
+        try:
+            staggered = field.staggered
+        except AttributeError:
+            return None
+        if not staggered or all(s == 0 for s in staggered):
+            return None
+        return tuple((d.spacing / 2) if s else 0
+                     for d, s in zip(self._gdims, staggered, strict=True))
+
     @memoized_meth
-    def _rdim(self, subdomain=None):
+    def _rdim(self, subdomain=None, shifts=None):
         # If the interpolation operation is limited to a SubDomain,
         # use the SubDimensions of that SubDomain
         if subdomain:
@@ -254,7 +270,7 @@ def _rdim(self, subdomain=None):
 
         # Make radius dimension conditional to avoid OOB
         rdims = []
-        pos = self.sfunction._position_map.values()
+        pos = self.sfunction._position_map(shifts=shifts).values()
 
         for (d, rd, p) in zip(gdims, self._cdim, pos, strict=True):
             # Add conditional to avoid OOB
@@ -300,27 +316,34 @@ def _augment_implicit_dims(self, implicit_dims, extras=None):
             idims = extra + as_tuple(implicit_dims) + self.sfunction.dimensions
         return tuple(idims)
 
-    def _coeff_temps(self, implicit_dims):
+    def _coeff_temps(self, implicit_dims, shifts=None):
         return []
 
-    def _positions(self, implicit_dims):
+    def _positions(self, implicit_dims, shifts=None):
         return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims)
-                for k, v in self.sfunction._position_map.items()]
+                for k, v in self.sfunction._position_map(shifts=shifts).items()]
 
-    def _interp_idx(self, variables, implicit_dims=None, subdomain=None):
+    def _interp_idx(self, variables, implicit_dims=None, subdomain=None,
+                    shifts=None):
         """
         Generate interpolation indices for the DiscreteFunctions in ``variables``.
+
+        ``shifts`` is a per-Dimension physical offset for the target field's
+        origin: it only affects the integer position symbol via the position
+        map (``pos = floor((c - o - shift)/h)``). The index substitution itself
+        is unchanged — any staggered offset in a field's own symbolic access is
+        absorbed by Devito's normal indexing.
         """
-        pos = self.sfunction._position_map.values()
+        pos = self.sfunction._position_map(shifts=shifts).values()
 
         # Temporaries for the position
-        temps = self._positions(implicit_dims)
+        temps = self._positions(implicit_dims, shifts=shifts)
 
         # Coefficient symbol expression
-        temps.extend(self._coeff_temps(implicit_dims))
+        temps.extend(self._coeff_temps(implicit_dims, shifts=shifts))
 
         # Substitution mapper for variables
-        mapper = self._rdim(subdomain=subdomain).getters
+        mapper = self._rdim(subdomain=subdomain, shifts=shifts).getters
 
         # Index substitution to make in variables
         subs = {
@@ -451,9 +474,11 @@ def _inject(self, field, expr, implicit_dims=None):
 
         subdomain = _extract_subdomain(fields)
 
-        # Derivatives must be evaluated before the introduction of indirect accesses
+        # Derivatives must be evaluated before the introduction of indirect
+        # accesses. Variables are sampled at their own grid location; the
+        # position map for the target field carries the staggering so the
+        # field's stencil neighbors land on the right indices.
         try:
-            # Expr must be evaluated ath the field's location.
             _exprs = tuple(e._eval_at(f).evaluate
                            for e, f in zip(exprs, fields, strict=True))
         except AttributeError:
@@ -464,22 +489,29 @@ def _inject(self, field, expr, implicit_dims=None):
 
         # Implicit dimensions
         implicit_dims = self._augment_implicit_dims(implicit_dims, variables)
+
+        # All fields in a single injection share the same staggering by
+        # construction (they are written together at the same indices), so
+        # derive shifts from the first field.
+        shifts = self._field_shifts(fields[0])
+
         # Move all temporaries inside inner loop to improve parallelism
-        # Can only be done for inject as interpolation need a temporary
-        # summing temp that wouldn't allow collapsing
+        # 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))
+                                              self._rdim(subdomain=subdomain,
+                                                         shifts=shifts))
 
         # List of indirection indices for all adjacent grid points
-        finterp = fields + as_tuple(variables)
-        idx_subs, temps = self._interp_idx(finterp, implicit_dims=implicit_dims,
-                                           subdomain=subdomain)
+        idx_subs, temps = self._interp_idx(list(fields) + variables,
+                                           implicit_dims=implicit_dims,
+                                           subdomain=subdomain, shifts=shifts)
 
-        # Substitute coordinate base symbols into the interpolation coefficients
         eqns = [Inc(_field.xreplace(idx_subs),
-                    (self._weights(subdomain=subdomain) * _expr).xreplace(idx_subs),
+                    (self._weights(subdomain=subdomain, shifts=shifts)
+                     * _expr).xreplace(idx_subs),
                     implicit_dims=implicit_dims)
-                for (_field, _expr) in zip(fields, _exprs, strict=True)]
+                for _field, _expr in zip(fields, _exprs, strict=True)]
 
         return temps + eqns
 
@@ -497,24 +529,27 @@ class LinearInterpolator(WeightedInterpolator):
     _name = 'linear'
 
     @memoized_meth
-    def _weights(self, subdomain=None):
-        rdim = self._rdim(subdomain=subdomain)
+    def _weights(self, subdomain=None, shifts=None):
+        rdim = self._rdim(subdomain=subdomain, shifts=shifts)
         c = [(1 - p) * (1 - r) + p * r
-             for (p, d, r) in zip(self._point_symbols, self._gdims, rdim, strict=True)]
+             for (p, d, r) in zip(self._point_symbols(shifts), self._gdims, rdim,
+                                  strict=True)]
         return Mul(*c)
 
-    @cached_property
-    def _point_symbols(self):
+    @memoized_meth
+    def _point_symbols(self, shifts=None):
         """Symbol for coordinate value in each Dimension of the point."""
         dtype = self.sfunction.coordinates.dtype
-        return DimensionTuple(*(Symbol(name=f'p{d}', dtype=dtype)
+        suffix = self.sfunction._shifts_suffix(shifts)
+        return DimensionTuple(*(Symbol(name=f'p{d}{suffix}', dtype=dtype)
                                 for d in self.grid.dimensions),
                               getters=self.grid.dimensions)
 
-    def _coeff_temps(self, implicit_dims):
+    def _coeff_temps(self, implicit_dims, shifts=None):
         # Positions
-        pmap = self.sfunction._position_map
-        poseq = [Eq(self._point_symbols[d], pos - floor(pos),
+        pmap = self.sfunction._position_map(shifts=shifts)
+        psyms = self._point_symbols(shifts)
+        poseq = [Eq(psyms[d], pos - floor(pos),
                     implicit_dims=implicit_dims)
                  for (d, pos) in zip(self._gdims, pmap.keys(), strict=True)]
         return poseq
@@ -533,23 +568,24 @@ class PrecomputedInterpolator(WeightedInterpolator):
 
     _name = 'precomp'
 
-    def _positions(self, implicit_dims):
+    def _positions(self, implicit_dims, shifts=None):
         if self.sfunction.gridpoints_data is None:
-            return super()._positions(implicit_dims)
+            return super()._positions(implicit_dims, shifts=shifts)
         else:
             # No position temp as we have directly the gridpoints
             return[Eq(p, k, implicit_dims=implicit_dims)
-                   for (k, p) in self.sfunction._position_map.items()]
+                   for (k, p) in self.sfunction._position_map(shifts=shifts).items()]
 
     @property
     def interpolation_coeffs(self):
         return self.sfunction.interpolation_coeffs
 
     @memoized_meth
-    def _weights(self, subdomain=None):
+    def _weights(self, subdomain=None, shifts=None):
         ddim, cdim = self.interpolation_coeffs.dimensions[1:]
         mappers = [{ddim: ri, cdim: rd-rd.parent.symbolic_min}
-                   for (ri, rd) in enumerate(self._rdim(subdomain=subdomain))]
+                   for (ri, rd) in enumerate(self._rdim(subdomain=subdomain,
+                                                        shifts=shifts))]
         return Mul(*[self.interpolation_coeffs.subs(mapper)
                      for mapper in mappers])
 
@@ -594,8 +630,8 @@ def interpolation_coeffs(self):
         return tuple(coeffs)
 
     @memoized_meth
-    def _weights(self, subdomain=None):
-        rdims = self._rdim(subdomain=subdomain)
+    def _weights(self, subdomain=None, shifts=None):
+        rdims = self._rdim(subdomain=subdomain, shifts=shifts)
         return Mul(*[
             w._subs(rd, rd-rd.parent.symbolic_min)
             for (rd, w) in zip(rdims, self.interpolation_coeffs, strict=True)

devito/types/sparse.py

@@ -366,11 +366,18 @@ def coordinates_data(self):
         except AttributeError:
             return None
 
-    @cached_property
-    def _pos_symbols(self):
-        return [Symbol(name=f'pos{d}', dtype=np.int32)
+    @memoized_meth
+    def _pos_symbols(self, shifts=None):
+        suffix = self._shifts_suffix(shifts)
+        return [Symbol(name=f'pos{d}{suffix}', dtype=np.int32)
                 for d in self.grid.dimensions]
 
+    @staticmethod
+    def _shifts_suffix(shifts):
+        if not shifts or all(s == 0 for s in shifts):
+            return ''
+        return '_s' + ''.join('1' if s != 0 else '0' for s in shifts)
+
     @cached_property
     def _point_increments(self):
         """Index increments in each Dimension for each point symbol."""
@@ -380,19 +387,25 @@ def _point_increments(self):
     def _point_support(self):
         return np.array(self._point_increments)
 
-    @cached_property
-    def _position_map(self):
+    @memoized_meth
+    def _position_map(self, shifts=None):
         """
-        Symbols map for the physical position of the sparse points relative to the grid
-        origin.
+        Symbols map for the physical position of the sparse points relative to the
+        origin of the target field. ``shifts`` is a tuple, one entry per grid
+        Dimension, of additional physical offsets to subtract from the sparse
+        coordinates (e.g. ``h_x/2`` for a field staggered in ``x``). If ``shifts``
+        is None, only the grid origin is subtracted.
         """
+        if shifts is None:
+            shifts = (0,) * len(self.grid.dimensions)
         return OrderedDict([
-            ((c - o)/d.spacing, p)
-            for p, c, d, o in zip(
-                self._pos_symbols,
+            ((c - o - s)/d.spacing, p)
+            for p, c, d, o, s in zip(
+                self._pos_symbols(shifts=shifts),
                 self._coordinate_symbols,
                 self.grid.dimensions,
                 self.grid.origin_symbols,
+                shifts,
                 strict=True
             )
         ])
@@ -419,7 +432,6 @@ def _cond_rdim(self, dim, cond):
         parent = self._crdim(dim)
         return ConditionalDimension(parent.name, parent, condition=cond, indirect=True)
 
-
     def _eval_at(self, func):
         return self
 
@@ -451,7 +463,7 @@ def guard(self, expr=None):
         conditions = {}
 
         # Position map and temporaries for it
-        pmap = self._position_map
+        pmap = self._position_map()
 
         # Temporaries for the position
         temps = self.interpolator._positions(self.dimensions)
@@ -713,9 +725,6 @@ def _dist_scatter(self, alias=None, data=None):
                 mapper.update(self._dist_subfunc_scatter(sf))
         return mapper
 
-    def _eval_at(self, func):
-        return self
-
     def _halo_exchange(self):
         # no-op for SparseFunctions
         return
@@ -1286,8 +1295,8 @@ def _coordinate_symbols(self):
             return tuple([self.coordinates._subs(d_dim, i)
                           for i in range(self.grid.dim)])
 
-    @cached_property
-    def _position_map(self):
+    @memoized_meth
+    def _position_map(self, shifts=None):
         """
         Symbol for each grid index according to the coordinates.
 
@@ -1307,12 +1316,12 @@ def _position_map(self):
                 (self.gridpoints._subs(ddim, di), p)
                 for (di, p) in zip(
                     range(self.grid.dim),
-                    self._pos_symbols,
+                    self._pos_symbols(shifts=shifts),
                     strict=True
                 )
             )
         else:
-            return super()._position_map
+            return super()._position_map(shifts=shifts)
 
 
 class PrecomputedSparseTimeFunction(AbstractSparseTimeFunction,

tests/test_interpolation.py

@@ -6,7 +6,7 @@
 
 from conftest import assert_structure
 from devito import (
-    DefaultDimension, Dimension, Eq, Function, Grid, MatrixSparseTimeFunction, NODE,
+    NODE, DefaultDimension, Dimension, Eq, Function, Grid, MatrixSparseTimeFunction,
     Operator, PrecomputedSparseFunction, PrecomputedSparseTimeFunction, Real,
     SparseFunction, SparseTimeFunction, SubDomain, TimeFunction, switchconfig
 )
@@ -485,21 +485,28 @@ def test_inject_staggered(self, stagg):
         op = Operator(expr)
 
         op()
-        # The expected value is 1 for the NODE case.
-        # For the staggered case, the center point is
-        # shifted left compared to the staggered function
         if stagg == 'NODE':
             assert np.isclose(a.data[5, 5, 5], 2, rtol=1e-6)
             # all other should be zero
             assert np.sum(a.data) == 2
-        elif len(as_tuple(staggered)) == 1:
-            from IPython import embed; embed()
-            # 
-            assert np.isclose(a.data[5, 5, 5], 1, rtol=1e-6)
-        elif len(as_tuple(staggered)) == 2:
-            expected = 0.25
-        elif len(as_tuple(staggered)) == 3:
-            expected = 0.125
+        else:
+            # Bottom corner since the source position is left of the staggered field
+            corner = [5, 5, 5] - np.array(a.staggered)
+            # Indices touched by the interpolation based on staggering
+            slices = [slice(corner[i], 6) for i in range(3)]
+            # Number of points for the interpolation.
+            # Single dim only interpolates between two points,
+            # so 2**(number of staggered dims) is the number of points.
+            npoints = 2**(np.sum(np.array(a.staggered, dtype=np.int32)))
+            # b value at the staggered field from b._eval_at(field)
+            b_val = (1 * (npoints - 1) + 2) / npoints
+            # Then source interpolation. Source at the center of the staggered field,
+            # so all points have the same weight 1/npoints.
+            interp_val = b_val / npoints
+            assert np.allclose(a.data[slices], interp_val, rtol=1e-6)
+            # All other should be zero so should sum to the interp_val * number of points.
+            # Use abs to make sure there is no +- cancellations
+            assert np.sum(np.abs(a.data)) == interp_val * 2**(sum(np.array(a.staggered)))
 
 # ---------------------------------------------------------------------------
 # Precomputed interpolation / injection
@@ -824,7 +831,7 @@ def test_point_symbol_types(self, dtype, expected):
         grid = Grid(shape=(11,))
         s = SparseFunction(name='src', npoint=1,
                            grid=grid, dtype=dtype)
-        point_symbol = s.interpolator._point_symbols[0]
+        point_symbol = s.interpolator._point_symbols()[0]
 
         assert point_symbol.dtype is expected