compiler: Support interpolation with user-provided implicit dims

devito/operations/interpolators.py

@@ -6,7 +6,7 @@
 
 from devito.logger import warning
 from devito.symbolics import retrieve_function_carriers, indexify, INT
-from devito.tools import powerset, flatten, prod
+from devito.tools import as_tuple, powerset, flatten, prod
 from devito.types import (ConditionalDimension, Dimension, DefaultDimension, Eq, Inc,
                           Evaluable, Symbol, SubFunction)
 
@@ -168,7 +168,8 @@ def _interpolation_coeffs(self):
         A = A.subs(reference_cell)
         return A.inv().T * p
 
-    def _interpolation_indices(self, variables, offset=0, field_offset=0):
+    def _interpolation_indices(self, variables, offset=0, field_offset=0,
+                               implicit_dims=None):
         """
         Generate interpolation indices for the DiscreteFunctions in ``variables``.
         """
@@ -192,21 +193,22 @@ def _interpolation_indices(self, variables, offset=0, field_offset=0):
             idx_subs.append(mapper)
 
         # Temporaries for the position
-        temps = [Eq(v, k, implicit_dims=self.sfunction.dimensions)
+        temps = [Eq(v, k, implicit_dims=implicit_dims)
                  for k, v in self.sfunction._position_map.items()]
         # Temporaries for the indirection dimensions
         temps.extend([Eq(v, k.subs(self.sfunction._position_map),
-                         implicit_dims=self.sfunction.dimensions)
+                         implicit_dims=implicit_dims)
                       for k, v in points.items()])
         # Temporaries for the coefficients
         temps.extend([Eq(p, c.subs(self.sfunction._position_map),
-                         implicit_dims=self.sfunction.dimensions)
+                         implicit_dims=implicit_dims)
                       for p, c in zip(self.sfunction._point_symbols,
                                       self.sfunction._coordinate_bases(field_offset))])
 
         return idx_subs, temps
 
-    def interpolate(self, expr, offset=0, increment=False, self_subs={}):
+    def interpolate(self, expr, offset=0, increment=False, self_subs={},
+                    implicit_dims=None):
         """
         Generate equations interpolating an arbitrary expression into ``self``.
 
@@ -218,7 +220,13 @@ def interpolate(self, expr, offset=0, increment=False, self_subs={}):
             Additional offset from the boundary.
         increment: bool, optional
             If True, generate increments (Inc) rather than assignments (Eq).
+        implicit_dims : Dimension or list of Dimension, optional
+            An ordered list of Dimensions that do not explicitly appear in the
+            interpolation expression, but that should be honored when constructing
+            the operator.
         """
+        implicit_dims = as_tuple(implicit_dims) + self.sfunction.dimensions
+
         def callback():
             # Derivatives must be evaluated before the introduction of indirect accesses
             try:
@@ -233,18 +241,18 @@ def callback():
             # TODO: handle each variable staggereing spearately
             field_offset = variables[0].origin
             # List of indirection indices for all adjacent grid points
-            idx_subs, temps = self._interpolation_indices(variables, offset,
-                                                          field_offset=field_offset)
+            idx_subs, temps = self._interpolation_indices(
+                variables, offset, field_offset=field_offset, implicit_dims=implicit_dims
+            )
 
             # Substitute coordinate base symbols into the interpolation coefficients
             args = [_expr.xreplace(v_sub) * b.xreplace(v_sub)
                     for b, v_sub in zip(self._interpolation_coeffs, idx_subs)]
 
             # Accumulate point-wise contributions into a temporary
             rhs = Symbol(name='sum', dtype=self.sfunction.dtype)
-            summands = [Eq(rhs, 0., implicit_dims=self.sfunction.dimensions)]
-            summands.extend([Inc(rhs, i, implicit_dims=self.sfunction.dimensions)
-                            for i in args])
+            summands = [Eq(rhs, 0., implicit_dims=implicit_dims)]
+            summands.extend([Inc(rhs, i, implicit_dims=implicit_dims) for i in args])
 
             # Write/Incr `self`
             lhs = self.sfunction.subs(self_subs)
@@ -254,7 +262,7 @@ def callback():
 
         return Interpolation(expr, offset, increment, self_subs, self, callback)
 
-    def inject(self, field, expr, offset=0):
+    def inject(self, field, expr, offset=0, implicit_dims=None):
         """
         Generate equations injecting an arbitrary expression into a field.
 
@@ -266,7 +274,13 @@ def inject(self, field, expr, offset=0):
             Injected expression.
         offset : int, optional
             Additional offset from the boundary.
+        implicit_dims : Dimension or list of Dimension, optional
+            An ordered list of Dimensions that do not explicitly appear in the
+            injection expression, but that should be honored when constructing
+            the operator.
         """
+        implicit_dims = as_tuple(implicit_dims) + self.sfunction.dimensions
+
         def callback():
             # Derivatives must be evaluated before the introduction of indirect accesses
             try:
@@ -280,12 +294,13 @@ def callback():
             # Need to get origin of the field in case it is staggered
             field_offset = field.origin
             # List of indirection indices for all adjacent grid points
-            idx_subs, temps = self._interpolation_indices(variables, offset,
-                                                          field_offset=field_offset)
+            idx_subs, temps = self._interpolation_indices(
+                variables, offset, field_offset=field_offset, implicit_dims=implicit_dims
+            )
 
             # Substitute coordinate base symbols into the interpolation coefficients
             eqns = [Inc(field.xreplace(vsub), _expr.xreplace(vsub) * b,
-                        implicit_dims=self.sfunction.dimensions)
+                        implicit_dims=implicit_dims)
                     for b, vsub in zip(self._interpolation_coeffs, idx_subs)]
 
             return temps + eqns

devito/types/sparse.py

@@ -848,7 +848,7 @@ def interpolate(self, expr, offset=0, u_t=None, p_t=None, increment=False):
                                                            increment=increment,
                                                            self_subs=subs)
 
-    def inject(self, field, expr, offset=0, u_t=None, p_t=None):
+    def inject(self, field, expr, offset=0, u_t=None, p_t=None, implicit_dims=None):
         """
         Generate equations injecting an arbitrary expression into a field.
 
@@ -864,14 +864,18 @@ def inject(self, field, expr, offset=0, u_t=None, p_t=None):
             Time index at which the interpolation is performed.
         p_t : expr-like, optional
             Time index at which the result of the interpolation is stored.
+        implicit_dims : Dimension or list of Dimension, optional
+            An ordered list of Dimensions that do not explicitly appear in the
+            injection expression, but that should be honored when constructing
+            the operator.
         """
         # 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(SparseTimeFunction, self).inject(field, expr, offset=offset)
+        return super().inject(field, expr, offset=offset, implicit_dims=implicit_dims)
 
     # Pickling support
     _pickle_kwargs = AbstractSparseTimeFunction._pickle_kwargs +\

tests/test_dimension.py

@@ -5,9 +5,10 @@
 import pytest
 
 from conftest import assert_blocking, skipif, opts_tiling
-from devito import (ConditionalDimension, Grid, Function, TimeFunction, SparseFunction,  # noqa
-                    Eq, Operator, Constant, Dimension, SubDimension, switchconfig,
-                    SubDomain, Lt, Le, Gt, Ge, Ne, Buffer)
+from devito import (ConditionalDimension, Grid, Function, TimeFunction,  # noqa
+                    SparseFunction, SparseTimeFunction, Eq, Operator, Constant,
+                    Dimension, SubDimension, switchconfig, SubDomain, Lt, Le,
+                    Gt, Ge, Ne, Buffer, sin)
 from devito.ir.iet import (Conditional, Expression, Iteration, FindNodes,
                            retrieve_iteration_tree)
 from devito.symbolics import indexify, retrieve_functions, IntDiv
@@ -1178,6 +1179,41 @@ def test_affiness(self):
         iterations = [i for i in FindNodes(Iteration).visit(op) if i.dim is not time]
         assert all(i.is_Affine for i in iterations)
 
+    def test_sparse_time_function(self):
+        nt = 20
+
+        shape = (21, 21, 21)
+        origin = (0., 0., 0.)
+        spacing = (1., 1., 1.)
+        extent = tuple([d * (s - 1) for s, d in zip(shape, spacing)])
+        grid = Grid(shape=shape, extent=extent, origin=origin)
+        time = grid.time_dim
+        x, y, z = grid.dimensions
+
+        p = TimeFunction(name="p", grid=grid, time_order=2, space_order=4, save=nt)
+
+        # Place source in the middle of the grid
+        src_coords = np.empty((1, len(shape)), dtype=np.float32)
+        src_coords[0, :] = [o + d * (s-1)//2 for o, d, s in zip(origin, spacing, shape)]
+        src = SparseTimeFunction(name='src', grid=grid, npoint=1, nt=nt)
+        src.data[:] = 1.
+        src.coordinates.data[:] = src_coords[:]
+
+        cd = ConditionalDimension(name="cd", parent=time,
+                                  condition=And(Ge(time, 1), Le(time, 10)))
+
+        src_term = src.inject(field=p.forward, expr=src*time, implicit_dims=cd)
+
+        op = Operator(src_term)
+
+        op.apply(time_m=1, time_M=nt-2, dt=1.0)
+
+        assert np.all(p.data[0] == 0)
+        # Note the endpoint of the range is 12 because we inject at p.forward
+        assert all(p.data[i].sum() == i - 1 for i in range(1, 12))
+        assert all(p.data[i, 10, 10, 10] == i - 1 for i in range(1, 12))
+        assert all(np.all(p.data[i] == 0) for i in range(12, 20))
+
 
 class TestMashup(object):