api: add option to apply interp indices to function coeffs at inject

Author: mloubout

devito/operations/interpolators.py

@@ -164,19 +164,21 @@ class Injection(UnevaluatedSparseOperation):
 
     __rargs__ = ('field', 'expr', 'implicit_dims') + UnevaluatedSparseOperation.__rargs__
 
-    def __new__(cls, field, expr, implicit_dims, interpolator):
+    def __new__(cls, field, expr, implicit_dims, interpolator, interp_expr=False):
         obj = super().__new__(cls, interpolator)
 
         # TODO: unused now, but will be necessary to compute the adjoint
         obj.field = field
         obj.expr = expr
         obj.implicit_dims = implicit_dims
+        obj.interp_expr = interp_expr
 
         return obj
 
     def operation(self, **kwargs):
         return self.interpolator._inject(expr=self.expr, field=self.field,
-                                         implicit_dims=self.implicit_dims)
+                                         implicit_dims=self.implicit_dims,
+                                         interp_expr=self.interp_expr)
 
     def __repr__(self):
         return f"Injection({repr(self.expr)} into {repr(self.field)})"
@@ -366,7 +368,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, implicit_dims=None, interp_expr=False):
         """
         Generate equations injecting an arbitrary expression into a field.
 
@@ -381,7 +383,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, implicit_dims, self, interp_expr=interp_expr)
 
     def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None):
         """
@@ -433,7 +435,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, implicit_dims=None, interp_expr=False):
         """
         Generate equations injecting an arbitrary expression into a field.
 
@@ -479,8 +481,10 @@ def _inject(self, field, expr, implicit_dims=None):
                                               self._rdim(subdomain=subdomain))
 
         # List of indirection indices for all adjacent grid points
-        idx_subs, temps = self._interp_idx(fields, implicit_dims=implicit_dims,
-                                           pos_only=variables, subdomain=subdomain)
+        finterp = fields + as_tuple(variables) if interp_expr else fields
+        pos_only = () if interp_expr else variables
+        idx_subs, temps = self._interp_idx(finterp, implicit_dims=implicit_dims,
+                                           pos_only=pos_only, subdomain=subdomain)
 
         # Substitute coordinate base symbols into the interpolation coefficients
         eqns = [Inc(_field.xreplace(idx_subs),

devito/types/sparse.py

@@ -1089,7 +1089,8 @@ 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, u_t=None, p_t=None, implicit_dims=None,
+               interp_expr=False):
         """
         Generate equations injecting an arbitrary expression into a field.
 
@@ -1114,7 +1115,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, implicit_dims=implicit_dims,
+                              interp_expr=interp_expr)
 
     @property
     def forward(self):

tests/test_interpolation.py

@@ -11,6 +11,7 @@
     SparseTimeFunction, SubDomain, TimeFunction, switchconfig
 )
 from devito.operations.interpolators import LinearInterpolator, SincInterpolator
+from devito.symbolics import retrieve_functions
 from examples.seismic import (
     AcquisitionGeometry, Receiver, RickerSource, TimeAxis, demo_model
 )
@@ -569,6 +570,37 @@ def test_inject_from_field(shape, coords, result, npoints=19):
     assert np.allclose(a.data[indices], result, rtol=1.e-5)
 
 
+def test_inject_interp_expr():
+    """
+    Test that the Function coefficient gets interpolated too.
+    """
+    coords = [(.05, .95), (.45, .45)]
+    a = unit_box(shape=(11, 11))
+    a.data[:] = 0.
+    p = points(a.grid, ranges=coords, npoints=19)
+    m = Function(name='m', grid=a.grid)
+    m.data_with_halo[:] = 1.
+
+    expr = p.inject(a, m, interp_expr=True)
+    op = Operator(expr)
+
+    op(a=a)
+
+    indices = [slice(4, 6, 1) for _ in coords]
+    indices[0] = slice(1, -1, 1)
+    assert np.allclose(a.data[indices], 1, rtol=1.e-5)
+
+    # Extract interp expr to check indices
+    e_expr = expr.evaluate
+    funcs = retrieve_functions(e_expr[-1])
+    assert m in {f.function for f in funcs}
+    # All funcs should have the same indices wit radius
+    # includint the coefficient m
+    indices = {f.indices for f in funcs}
+    assert len(indices) == 1
+    assert str(indices.pop()) == '(rp_pointsx + posx, rp_pointsy + posy)'
+
+
 @pytest.mark.parametrize('shape', [(50, 50, 50)])
 def test_position(shape):
     t0 = 0.0  # Start time