compiler: source injection to nested field

devito/operations/interpolators.py

@@ -292,8 +292,17 @@ def callback():
 
             variables = list(retrieve_function_carriers(_expr)) + [field]
 
-            # Need to get origin of the field in case it is staggered
-            field_offset = field.origin
+            try:
+                # Search if the injection field contains nested fields. In case there
+                # are, assert they have the same origin and use them for field_offset
+                fields = retrieve_function_carriers(field.indices)
+                origin = [f.origin for f in fields]
+                assert all(o == origin[0] for o in origin)
+                field_offset = origin[0]
+            except:
+                # 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, implicit_dims=implicit_dims

devito/types/sparse.py

@@ -572,6 +572,16 @@ def gridpoints(self):
             np.floor(self.coordinates.data._local - self.grid.origin) / self.grid.spacing
         ).astype(int)
 
+    @property
+    def gridpoints_all(self):
+        if self.coordinates._data is None:
+            raise ValueError("No coordinates attached to this SparseFunction")
+
+        arr = np.moveaxis(self._support, -1, 0)
+        arr = arr.reshape(np.prod(arr.shape[:-1]), arr.shape[-1])
+        arr = np.unique(arr, axis=0)
+        return arr
+
     def guard(self, expr=None, offset=0):
         """
         Generate guarded expressions, that is expressions that are evaluated

tests/test_interpolation.py

@@ -9,7 +9,7 @@
                     PrecomputedSparseFunction, PrecomputedSparseTimeFunction,
                     MatrixSparseTimeFunction)
 from examples.seismic import (demo_model, TimeAxis, RickerSource, Receiver,
-                              AcquisitionGeometry)
+                              AcquisitionGeometry, Model)
 from examples.seismic.acoustic import AcousticWaveSolver
 import scipy.sparse
 
@@ -625,3 +625,78 @@ class SparseFirst(SparseFunction):
     op(time_M=10)
     expected = 10*11/2  # n (n+1)/2
     assert np.allclose(s.data, expected)
+
+
+@pytest.mark.parametrize('inj', ('s_id', '1 + s_id', 's_id[0, 0, s_id]'))
+@pytest.mark.parametrize('shape', [(50, 50, 50)])
+@pytest.mark.parametrize('so', (2, 4, 8))
+@pytest.mark.parametrize('tn', (20, 40, 60))
+def test_decompose_src_to_aligned(shape, so, tn, inj):
+    """ Test decomposition of non-aligned source wavelets to equivalent
+        aligned to grid points source wavelets
+    """
+
+    spacing = (10., 10., 10)
+    origin = (0., 0., 0.)
+
+    # Initialize v field
+    v = np.empty(shape, dtype=np.float32)
+    v[:, :, :int(shape[2]/2)] = 2
+    v[:, :, int(shape[2]/2):] = 1
+
+    # Construct model
+    model = Model(vp=v, origin=origin, shape=shape, spacing=spacing, space_order=so)
+
+    t0 = 0  # Simulation starts a t=0
+    dt = 1  # model.critical_dt  # Time step from model grid spacing
+    tn = tn
+    time_range = TimeAxis(start=t0, stop=tn, step=dt)
+    f0 = 0.010  # Source peak frequency is 10Hz (0.010 kHz)
+    src = RickerSource(name='src', grid=model.grid, f0=f0,
+                       npoint=9, time_range=time_range)
+
+    # First, position source centrally in all dimensions, then set depth
+    stx = 0.125
+    ste = 0.9
+    stepx = (ste-stx)/int(np.sqrt(src.npoint))
+
+    # Uniform x, y source spread
+    src.coordinates.data[:, :2] = \
+        np.array(np.meshgrid(np.arange(stx, ste,
+                 stepx), np.arange(stx, ste, stepx))).T.reshape(-1, 2) \
+        * np.array(model.domain_size[:1])
+
+    src.coordinates.data[:, -1] = 20  # Depth is 20m
+
+    # Source ID function to hold unique id for each point affected
+    s_id = Function(name='s_id', shape=model.grid.shape, dimensions=model.grid.dimensions,
+                    space_order=0, dtype=np.int32)
+
+    # Get positions affected by sparse operator
+    arr = src.gridpoints_all
+    nzinds = (arr[:, 0], arr[:, 1], arr[:, 2])
+    s_id.data[nzinds] = tuple(np.arange(len(nzinds[0])))
+
+    # Helper dimension to schedule loops of different sizes together
+    id_dim = Dimension(name='id_dim')
+
+    time = model.grid.time_dim
+    save_src = TimeFunction(name='save_src', shape=(src.shape[0], len(arr)),
+                            dimensions=(time, id_dim))
+
+    inj = eval(inj)
+    save_src_term = src.inject(field=save_src[time, inj],
+                               expr=src * dt**2 / model.m)
+
+    op1 = Operator(save_src_term)
+    op1.apply()
+
+    # Assert that first, last as well as other indices are as expected
+    assert(s_id.data[nzinds[0][0], nzinds[1][0], nzinds[2][0]] == 0)
+    assert(s_id.data[nzinds[0][-1], nzinds[1][-1], nzinds[2][-1]] == len(nzinds[0])-1)
+    assert(s_id.data[nzinds[0][len(nzinds[0])-1], nzinds[1][len(nzinds[0])-1],
+           nzinds[2][len(nzinds[0])-1]] == len(nzinds[0])-1)
+
+    # Assert that first, last as well as other indices are as expected
+    assert (src.shape[0] == save_src.shape[0])
+    assert (8*src.shape[1] == save_src.shape[1])