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,7 @@
 from abc import ABC, abstractmethod
+from contextlib import suppress
 from functools import cached_property, wraps
+from itertools import groupby
 
 import numpy as np
 import sympy
@@ -67,12 +69,9 @@ def _extract_subdomain(variables):
     """
     sdms = set()
     for v in variables:
-        try:
+        with suppress(AttributeError):
             if v.grid.is_SubDomain:
                 sdms.add(v.grid)
-        except AttributeError:
-            # Variable not on a grid (Indexed for example)
-            pass
 
     if len(sdms) > 1:
         raise NotImplementedError("Sparse operation on multiple Functions defined on"
@@ -245,19 +244,17 @@ def _cdim(self):
 
     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
+        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
+        staggered = field.staggered
+        if not staggered or staggered.on_node:
+            return ()
         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):
@@ -295,12 +292,10 @@ def _augment_implicit_dims(self, implicit_dims, extras=None):
             # dimensions of that SubDomain from any extra dimensions found
             edims = []
             for v in extras:
-                try:
+                with suppress(AttributeError):
                     if v.grid.is_SubDomain:
                         edims.extend([d for d in v.grid.dimensions
                                       if d.is_Sub and d.root in self._gdims])
-                except AttributeError:
-                    pass
 
             gdims = filter_ordered(edims + list(self._gdims))
             extra = filter_ordered([i for v in extras for i in v.dimensions
@@ -326,11 +321,11 @@ def _positions(self, implicit_dims, shifts=None):
     def _interp_idx(self, variables, implicit_dims=None, subdomain=None,
                     shifts=None):
         """
-        Generate interpolation indices for the DiscreteFunctions in ``variables``.
+        Generate interpolation indices for the DiscreteFunctions in `variables`.
 
-        ``shifts`` is a per-Dimension physical offset for the target field's
+        `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
+        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.
         """
@@ -360,7 +355,7 @@ def _interp_idx(self, variables, implicit_dims=None, subdomain=None,
     @check_coords
     def interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None):
         """
-        Generate equations interpolating an arbitrary expression into ``self``.
+        Generate equations interpolating an arbitrary expression into `self`.
 
         Parameters
         ----------
@@ -398,7 +393,7 @@ def inject(self, field, expr, implicit_dims=None):
 
     def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None):
         """
-        Generate equations interpolating an arbitrary expression into ``self``.
+        Generate equations interpolating an arbitrary expression into `self`.
 
         Parameters
         ----------
@@ -412,16 +407,13 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None
             the operator.
         """
         # Derivatives must be evaluated before the introduction of indirect accesses
-        try:
-            _expr = expr._eval_at(self.sfunction).evaluate
-        except AttributeError:
-            # E.g., a generic SymPy expression or a number
-            _expr = expr
+        with suppress(AttributeError):
+            expr = expr._eval_at(self.sfunction).evaluate
 
         if self_subs is None:
             self_subs = {}
 
-        variables = list(retrieve_function_carriers(_expr))
+        variables = list(retrieve_function_carriers(expr))
         subdomain = _extract_subdomain(variables)
 
         # Implicit dimensions
@@ -436,7 +428,7 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None
         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),
+        summands.extend([Inc(rhs, (weights * expr).xreplace(idx_subs),
                              implicit_dims=implicit_dims)])
 
         # Write/Incr `self`
@@ -478,42 +470,44 @@ def _inject(self, field, expr, implicit_dims=None):
         # 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:
-            _exprs = tuple(e._eval_at(f).evaluate
-                           for e, f in zip(exprs, fields, strict=True))
-        except AttributeError:
-            # 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
+        with suppress(AttributeError):
+            exprs = tuple(e._eval_at(f).evaluate
+                          for e, f in zip(exprs, fields, strict=True))
+
+        eqns = []
+        temps = []
+        # 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(zip(fields, exprs, strict=True), lambda f: f[0].staggered):
+            g_fields, g_exprs = zip(*g, strict=True)
+            variables = list(v for e in g_exprs for v in retrieve_function_carriers(e))
+
+            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(list(g_fields) + variables,
+                                                implicit_dims=implicit_dims,
+                                                subdomain=subdomain, shifts=shifts)
+
+            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)
+                         for f, e in zip(g_fields, g_exprs, strict=True)])
+
+        return filter_ordered(temps) + eqns
 
 
 class LinearInterpolator(WeightedInterpolator):
@@ -540,10 +534,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/basic.py

@@ -960,12 +960,12 @@ def indices(self):
         """The indices of the object."""
         return DimensionTuple(*self.args, getters=self.dimensions)
 
-    @property
+    @cached_property
     def indices_ref(self):
         """The reference indices of the object (indices at first creation)."""
         return DimensionTuple(*self.function.indices, getters=self.dimensions)
 
-    @property
+    @cached_property
     def origin(self):
         """
         Origin of the AbstractFunction in term of Dimension

devito/types/sparse.py

@@ -257,16 +257,17 @@ def sparse_position(self):
     def _sparse_dim(self):
         return self.dimensions[self.sparse_position]
 
-    @property
+    @cached_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):
         return {}
 
-    @property
+    @cached_property
     def origin(self):
         return DimensionTuple(*[0]*len(self.dimensions),
                               getters=self.dimensions)
@@ -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):
@@ -396,8 +395,7 @@ def _position_map(self, shifts=None):
         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)
+        shifts = shifts or (0,) * len(self.grid.dimensions)
         return OrderedDict([
             ((c - o - s)/d.spacing, p)
             for p, c, d, o, s in zip(

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