api: fix pickling of custom sparse functions with more than one dim

Author: mloubout

devito/types/dense.py

@@ -63,7 +63,8 @@ class DiscreteFunction(AbstractFunction, ArgProvider, Differentiable):
     The type of the underlying data object.
     """
 
-    __rkwargs__ = AbstractFunction.__rkwargs__ + ('staggered', 'coefficients')
+    __rkwargs__ = AbstractFunction.__rkwargs__ + \
+        ('shape_global', 'staggered', 'coefficients')
 
     def __init_finalize__(self, *args, function=None, **kwargs):
         # Now that *all* __X_setup__ hooks have been called, we can let the
@@ -1009,7 +1010,7 @@ class Function(DiscreteFunction):
     is_autopaddable = True
 
     __rkwargs__ = (DiscreteFunction.__rkwargs__ +
-                   ('space_order', 'shape_global', 'dimensions'))
+                   ('space_order', 'dimensions'))
 
     def _cache_meta(self):
         # Attach additional metadata to self's cache entry

devito/types/sparse.py

@@ -95,12 +95,19 @@ def __shape_setup__(cls, **kwargs):
         # A Grid must have been provided
         if grid is None:
             raise TypeError('Need `grid` argument')
-        shape = kwargs.get('shape')
+        shape = kwargs.get('shape', kwargs.get('shape_global'))
         dimensions = kwargs.get('dimensions')
         npoint = kwargs.get('npoint', kwargs.get('npoint_global'))
         glb_npoint = SparseDistributor.decompose(npoint, grid.distributor)
+        # Plain SparseFunction construction with npoint.
         if shape is None:
             loc_shape = (glb_npoint[grid.distributor.myrank],)
+        # No dimensions is only possible through rebuild, the shape is from
+        # the existing function
+        elif dimensions is None:
+            loc_shape = list(shape)
+            # For safety, ensure the distributed sparse dimension is correct
+            loc_shape[cls._sparse_position] = glb_npoint[grid.distributor.myrank]
         else:
             loc_shape = []
             assert len(dimensions) == len(shape)
@@ -111,6 +118,7 @@ def __shape_setup__(cls, **kwargs):
                     loc_shape.append(grid.size_map[d].loc)
                 else:
                     loc_shape.append(s)
+
         return tuple(loc_shape)
 
     def __fd_setup__(self):
@@ -733,16 +741,19 @@ def time_dim(self):
 
     @classmethod
     def __shape_setup__(cls, **kwargs):
-        shape = kwargs.get('shape')
-        if shape is None:
-            nt = kwargs.get('nt')
-            if not isinstance(nt, int):
-                raise TypeError('Need `nt` int argument')
-            if nt <= 0:
-                raise ValueError('`nt` must be > 0')
-
-            shape = list(AbstractSparseFunction.__shape_setup__(**kwargs))
-            shape.insert(cls._time_position, nt)
+        shape = list(AbstractSparseFunction.__shape_setup__(**kwargs))
+        dimensions = as_tuple(kwargs.get('dimensions'))
+        if dimensions is None or len(shape) == len(dimensions):
+            # Shape has already been setup, for example via rebuild
+            return tuple(shape)
+
+        nt = kwargs.get('nt')
+        if not isinstance(nt, int):
+            raise TypeError('Need `nt` int argument')
+        if nt <= 0:
+            raise ValueError('`nt` must be > 0')
+
+        shape.insert(cls._time_position, nt)
 
         return tuple(shape)
 

tests/test_pickle.py

@@ -19,7 +19,7 @@
 from devito.types import (Array, CustomDimension, Symbol as dSymbol, Scalar,
                           PointerArray, Lock, PThreadArray, SharedData, Timer,
                           DeviceID, NPThreads, ThreadID, TempFunction, Indirection,
-                          FIndexed, ComponentAccess)
+                          FIndexed, ComponentAccess, DefaultDimension)
 from devito.types.basic import BoundSymbol, AbstractSymbol
 from devito.tools import EnrichedTuple
 from devito.symbolics import (IntDiv, ListInitializer, FieldFromPointer,
@@ -29,6 +29,11 @@
                               TimeAxis, RickerSource, Receiver)
 
 
+class SparseFirst(SparseFunction):
+
+    _sparse_position = 0
+
+
 class SD(SubDomain):
     name = 'sd'
 
@@ -181,6 +186,32 @@ def test_precomputed_sparse_function(self, mode, pickle):
         assert sf.dtype == new_sf.dtype
         assert sf.npoint == new_sf.npoint == 3
 
+    def test_sparse_first(self, pickle):
+
+        dr = Dimension("cd")
+        ds = DefaultDimension("ps", default_value=3)
+        grid = Grid((11, 11))
+        sf = SparseFirst(name="s", grid=grid, npoint=2,
+                         dimensions=(dr, ds), shape=(2, 3),
+                         coordinates=[[.5, .5], [.2, .2]])
+        sf.data[0] = 1.
+
+        pkl_sf = pickle.dumps(sf)
+        new_sf = pickle.loads(pkl_sf)
+
+        # .data is initialized, so it should have been pickled too
+        assert np.all(sf.data[0] == 1.)
+        assert np.all(new_sf.data[0] == 1.)
+        assert new_sf.interpolation == sf.interpolation
+
+        # coordinates should also have been pickled
+        assert np.all(sf.coordinates.data == new_sf.coordinates.data)
+
+        assert sf.space_order == new_sf.space_order
+        assert sf.dtype == new_sf.dtype
+        assert sf.npoint == new_sf.npoint
+        assert sf.shape == new_sf.shape
+
     def test_alias_sparse_function(self, pickle):
         grid = Grid(shape=(3,))
         sf = SparseFunction(name='sf', grid=grid, npoint=3, space_order=2,