api: fix inject handling of different staggering

Author: mloubout

devito/finite_differences/derivative.py

@@ -489,6 +489,9 @@ def _eval_at(self, func):
         # 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
 
         # Check if x0's keys come from a DerivedDimension
         x0 = func.indices_ref.getters

devito/operations/interpolators.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from functools import cached_property, wraps
+from itertools import groupby
 
 import numpy as np
 import sympy
@@ -257,7 +258,8 @@ def _field_shifts(self, field):
         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))
+                     for d, s in zip(field.dimensions, staggered, strict=True)
+                     if d.is_Space)
 
     @memoized_meth
     def _rdim(self, subdomain=None, shifts=None):
@@ -485,35 +487,43 @@ def _inject(self, field, expr, implicit_dims=None):
             # E.g., a generic SymPy expression or a number
             _exprs = exprs
 
-        variables = list(v for e in _exprs for v in retrieve_function_carriers(e))
-
-        # 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 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))
-
-        # List of indirection indices for all adjacent grid points
-        idx_subs, temps = self._interp_idx(list(fields) + variables,
-                                           implicit_dims=implicit_dims,
-                                           subdomain=subdomain, shifts=shifts)
-
-        eqns = [Inc(_field.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)]
-
-        return temps + eqns
+        eqns = []
+        temps = []
+        pairs = zip(fields, _exprs, strict=True)
+        # We need to create one set of equations (temps and and coeffs) per staggering
+        # field in which we inject as the reference index depends on the field's origin
+        for _, g in groupby(pairs, lambda f: f[0].staggered):
+            g = list(g)
+            g_fields = [f for f, _ in g]
+            g_exprs = [e for _, e in g]
+            variables = list(v for e in g_exprs for v in retrieve_function_carriers(e))
+
+            # 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(g_fields[0])
+
+            # 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))
+
+            # List of indirection indices for all adjacent grid points
+            idx_subs, _temps = self._interp_idx(g_fields + variables,
+                                                implicit_dims=implicit_dims,
+                                                subdomain=subdomain, shifts=shifts)
+            w = self._weights(subdomain=subdomain, shifts=shifts)
+            temps.extend(_temps)
+            eqns.extend([Inc(_field.xreplace(idx_subs), (w * _expr).xreplace(idx_subs),
+                             implicit_dims=implicit_dims)
+                         for _field, _expr in zip(g_fields, g_exprs, strict=True)])
+
+        return filter_ordered(temps) + eqns
 
 
 class LinearInterpolator(WeightedInterpolator):
@@ -540,10 +550,13 @@ def _weights(self, subdomain=None, shifts=None):
     def _point_symbols(self, shifts=None):
         """Symbol for coordinate value in each Dimension of the point."""
         dtype = self.sfunction.coordinates.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)
+        symbols = []
+        for d in self.grid.dimensions:
+            if shifts and shifts[self.grid.dimensions.index(d)] != 0:
+                symbols.append(Symbol(name=f'p{d}_s1', dtype=dtype))
+            else:
+                symbols.append(Symbol(name=f'p{d}', dtype=dtype))
+        return DimensionTuple(*symbols, getters=self.grid.dimensions)
 
     def _coeff_temps(self, implicit_dims, shifts=None):
         # Positions

devito/types/sparse.py

@@ -259,8 +259,9 @@ def _sparse_dim(self):
 
     @property
     def indices_ref(self):
-        return DimensionTuple(*self.grid.dimensions,
-                              getters=self.grid.dimensions)
+        getters = (*self.dimensions, *self.grid.dimensions)
+        indices = (*self.dimensions, *self.grid.dimensions)
+        return DimensionTuple(*indices, getters=getters)
 
     @property
     def _grid_map(self):
@@ -368,15 +369,13 @@ def coordinates_data(self):
 
     @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)
+        symbols = []
+        for d in self.grid.dimensions:
+            if shifts and shifts[self.grid.dimensions.index(d)] != 0:
+                symbols.append(Symbol(name=f'pos{d}_s1', dtype=np.int32))
+            else:
+                symbols.append(Symbol(name=f'pos{d}', dtype=np.int32))
+        return symbols
 
     @cached_property
     def _point_increments(self):

examples/seismic/tutorials/06_elastic_varying_parameters.ipynb

@@ -604,8 +604,8 @@
    "outputs": [],
    "source": [
     "assert np.isclose(norm(rec), 23.504, atol=0, rtol=1e-3)\n",
-    "assert np.isclose(norm(rec2), 2.425, atol=0, rtol=1e-3)\n",
-    "assert np.isclose(norm(rec3), 2.889, atol=0, rtol=1e-3)"
+    "assert np.isclose(norm(rec2), 2.4298, atol=0, rtol=1e-3)\n",
+    "assert np.isclose(norm(rec3), 2.7481, atol=0, rtol=1e-3)"
    ]
   },
   {
@@ -838,8 +838,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "assert np.isclose(norm(rec2), .3250, atol=0, rtol=1e-3)\n",
-    "assert np.isclose(norm(rec3), .26745, atol=0, rtol=1e-3)"
+    "assert np.isclose(norm(rec2), .30388, atol=0, rtol=1e-3)\n",
+    "assert np.isclose(norm(rec3), .26633, atol=0, rtol=1e-3)"
    ]
   },
   {
@@ -1102,9 +1102,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "assert np.isclose(norm(rec), 31.23, atol=0, rtol=1e-3)\n",
-    "assert np.isclose(norm(rec2), 3.5482, atol=0, rtol=1e-3)\n",
-    "assert np.isclose(norm(rec3), 4.7007, atol=0, rtol=1e-3)"
+    "assert np.isclose(norm(rec), 29.538, atol=0, rtol=1e-3)\n",
+    "assert np.isclose(norm(rec2), 1.9116, atol=0, rtol=1e-3)\n",
+    "assert np.isclose(norm(rec3), 3.4919, atol=0, rtol=1e-3)"
    ]
   }
  ],

tests/test_dimension.py

@@ -1274,10 +1274,10 @@ def test_no_index_sparse(self):
 
         radius = 1
         indices = [(INT(floor(i)), INT(floor(i))+radius)
-                   for i in sf._position_map]
+                   for i in sf._position_map()]
         bounds = [i.symbolic_size - radius for i in grid.dimensions]
 
-        eqs = [Eq(p, v) for (v, p) in sf._position_map.items()]
+        eqs = [Eq(p, v) for (v, p) in sf._position_map().items()]
         for e, i in enumerate(product(*indices)):
             args = [j > 0 for j in i]
             args.extend([j < k for j, k in zip(i, bounds, strict=True)])

tests/test_interpolation.py

@@ -8,7 +8,8 @@
 from devito import (
     NODE, DefaultDimension, Dimension, Eq, Function, Grid, MatrixSparseTimeFunction,
     Operator, PrecomputedSparseFunction, PrecomputedSparseTimeFunction, Real,
-    SparseFunction, SparseTimeFunction, SubDomain, TimeFunction, switchconfig
+    SparseFunction, SparseTimeFunction, SubDomain, TimeFunction, VectorFunction,
+    switchconfig
 )
 from devito.operations.interpolators import LinearInterpolator, SincInterpolator
 from devito.tools import as_tuple
@@ -472,7 +473,7 @@ def test_inject_staggered(self, stagg):
                      staggered=staggered)
         a.data.fill(0)
 
-        b = Function(name='b', grid=a.grid, space_order=2,
+        b = Function(name='b', grid=a.grid, space_order=8,
                      staggered=NODE)
         b.data.fill(1)
         b.data[5, 5, 5] = 2
@@ -508,6 +509,26 @@ def test_inject_staggered(self, stagg):
             # Use abs to make sure there is no +- cancellations
             assert np.sum(np.abs(a.data)) == interp_val * 2**(sum(np.array(a.staggered)))
 
+    def test_inject_staggered_mixed(self):
+        grid = Grid((11, 11, 11))
+        v = VectorFunction(name='v', grid=grid, space_order=2)
+        b = Function(name='b', grid=grid, space_order=2, staggered=NODE)
+        p = SparseFunction(name="p", grid=grid, nt=10, npoint=1)
+
+        eq = p.inject(v, expr=b * p).evaluate
+
+        # We should have
+        # - 3 injection equations v_x, v_y, v_z
+        # The standard 6 on node temps posx, posy, posz, px, py, pz
+        # 2 temps for the staggered in x vx posz_s1, px_s1
+        # 2 temps for the staggered in y vy posz_s1, py_s1
+        # 2 temps for the staggered in z vz posz_s1, pz_s1
+        assert len(eq) == 3 + 6 + 2 + 2 + 2
+
+        op = Operator(eq)
+        # 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')
+
 # ---------------------------------------------------------------------------
 # Precomputed interpolation / injection
 # ---------------------------------------------------------------------------