compiler: Tweak SparseFunction reconstruction

Author: FabioLuporini

devito/types/sparse.py

@@ -49,7 +49,8 @@ class AbstractSparseFunction(DiscreteFunction):
     _sub_functions = ()
     """SubFunctions encapsulated within this AbstractSparseFunction."""
 
-    __rkwargs__ = DiscreteFunction.__rkwargs__ + ('npoint_global', 'space_order')
+    __rkwargs__ = (DiscreteFunction.__rkwargs__ +
+                   ('dimensions', 'npoint_global', 'space_order'))
 
     def __init_finalize__(self, *args, **kwargs):
         super().__init_finalize__(*args, **kwargs)
@@ -133,14 +134,17 @@ def __subfunc_setup__(self, key, suffix, dtype=None):
             shape = (self.npoint, self.grid.dim)
 
         # Check if already a SubFunction
+        d = self.indices[self._sparse_position]
         if isinstance(key, SubFunction):
-            # Need to rebuild so the dimensions match the parent SparseFunction
-            indices = (self.indices[self._sparse_position], *key.indices[1:])
-            return key._rebuild(*indices, name=name, shape=shape,
-                                alias=self.alias, halo=None)
-        elif key is not None and not isinstance(key, Iterable):
-            raise ValueError("`%s` must be either SubFunction "
-                             "or iterable (e.g., list, np.ndarray)" % key)
+            if d in key.dimensions and not self.alias:
+                # From a reconstruction which leaves `dimensions` intact
+                return key
+            else:
+                # Need to rebuild so the dimensions match the parent
+                # SparseFunction, for example we end up here via `.subs(d, new_d)`
+                indices = (d, *key.indices[1:])
+                return key._rebuild(*indices, name=name, shape=shape,
+                                    alias=self.alias, halo=None)
 
         if key is None:
             # Fallback to default behaviour

tests/test_mpi.py

@@ -2633,14 +2633,17 @@ def run_adjoint_F(self, nd):
         solver = acoustic_setup(shape=shape, spacing=[15. for _ in shape],
                                 tn=tn, space_order=so, nrec=nrec,
                                 preset='layers-isotropic', dtype=np.float64)
+
         # Run forward operator
-        rec, u, _ = solver.forward()
+        src = solver.geometry.src
+        rec, u, _ = solver.forward(src=src)
 
         assert np.isclose(norm(u) / Eu, 1.0)
         assert np.isclose(norm(rec) / Erec, 1.0)
 
         # Run adjoint operator
-        srca, v, _ = solver.adjoint(rec=rec)
+        srca = src.func(name='srca')
+        srca, v, _ = solver.adjoint(srca=srca, rec=rec)
 
         assert np.isclose(norm(v) / Ev, 1.0)
         assert np.isclose(norm(srca) / Esrca, 1.0)

tests/test_sparse.py

@@ -417,14 +417,13 @@ def test_rebuild(self, sptype):
                 assert getattr(sp, subf).name.startswith("s_")
 
         # Rebuild with different name, this should drop the function
-        # and create new data
+        # and create new data, while the coordinates and more generally all
+        # SubFunctions remain the same
         sp2 = sp._rebuild(name="sr")
-
-        # Check new subfunction
         for subf in sp2._sub_functions:
             if getattr(sp2, subf) is not None:
-                assert getattr(sp2, subf).name.startswith("sr_")
-                assert np.all(getattr(sp2, subf).data == 0)
+                assert getattr(sp2, subf).name.startswith("s_")
+                assert np.all(getattr(sp2, subf).data == getattr(sp, subf).data)
 
         # Rebuild with different name as an alias
         sp2 = sp._rebuild(name="sr2", alias=True)
@@ -433,6 +432,14 @@ def test_rebuild(self, sptype):
                 assert getattr(sp2, subf).name.startswith("sr2_")
                 assert getattr(sp2, subf).data is None
 
+        # Rebuild with different name and dimensions. This is expected to recreate
+        # the SubFunctions as well
+        sp2 = sp._rebuild(name="sr3", dimensions=None)
+        for subf in sp2._sub_functions:
+            if getattr(sp2, subf) is not None:
+                assert getattr(sp2, subf).name.startswith("sr3_")
+                assert np.all(getattr(sp2, subf).data == 0)
+
     @pytest.mark.parametrize('sptype', _sptypes)
     def test_subs(self, sptype):
         grid = Grid((3, 3, 3))