updates map API for consistency across Images and Series

thunder/base.py

@@ -402,7 +402,7 @@ def filter(self, func):
 
         return self._constructor(filtered, labels=newlabels).__finalize__(self, noprop=('labels',))
 
-    def _map(self, func, axis=(0,), value_shape=None, dtype=None, with_keys=False):
+    def map(self, func, value_shape=None, dtype=None, with_keys=False):
         """
         Apply an array -> array function across an axis.
 
@@ -429,6 +429,8 @@ def _map(self, func, axis=(0,), value_shape=None, dtype=None, with_keys=False):
         with_keys : bool, optional, default=False
             Include keys as an argument to the function
         """
+        axis = self.baseaxes
+
         if self.mode == 'local':
             axes = sorted(tupleize(axis))
             key_shape = [self.shape[axis] for axis in axes]

thunder/blocks/blocks.py

@@ -49,18 +49,18 @@ def collect_blocks(self):
         if self.mode == 'local':
             return self.values.values.flatten().tolist()
 
-    def map(self, func, dims=None, dtype=None):
+    def map(self, func, value_shape=None, dtype=None):
         """
         Apply an array -> array function to each block
         """
-        mapped = self.values.map(func, value_shape=dims, dtype=dtype)
+        mapped = self.values.map(func, value_shape=value_shape, dtype=dtype)
         return self._constructor(mapped).__finalize__(self, noprop=('dtype',))
 
     def map_generic(self, func):
         """
         Apply an arbitrary array -> object function to each blocks.
         """
-        return self.values.map_generic(func)[0] 
+        return self.values.map_generic(func)[0]
 
     def first(self):
         """

thunder/images/images.py

@@ -185,26 +185,6 @@ def sample(self, nsamples=100, seed=None):
 
         return self._constructor(result)
 
-    def map(self, func, dims=None, dtype=None, with_keys=False):
-        """
-        Map an array -> array function over each image.
-
-        Parameters
-        ----------
-        func : function
-            The function to apply in the map.
-
-        dims : tuple, optional, default = None
-            If known, the dimensions of the data following function evaluation.
-
-        dtype : numpy.dtype, optional, default = None
-            If known, the type of the data following function evaluation.
-
-        with_keys : boolean, optional, default = False
-            If true, function should be of both tuple indices and values.
-        """
-        return self._map(func, axis=0, value_shape=dims, dtype=dtype, with_keys=with_keys)
-
     def reduce(self, func):
         """
         Reduce a function over images.
@@ -274,7 +254,7 @@ def max_projection(self, axis=2):
 
         newdims = list(self.dims)
         del newdims[axis]
-        return self.map(lambda x: amax(x, axis), dims=newdims)
+        return self.map(lambda x: amax(x, axis), value_shape=newdims)
 
     def max_min_projection(self, axis=2):
         """
@@ -293,7 +273,7 @@ def max_min_projection(self, axis=2):
 
         newdims = list(self.dims)
         del newdims[axis]
-        return self.map(lambda x: amax(x, axis) + amin(x, axis), dims=newdims)
+        return self.map(lambda x: amax(x, axis) + amin(x, axis), value_shape=newdims)
 
     def subsample(self, factor):
         """
@@ -321,7 +301,7 @@ def roundup(a, b):
         slices = [slice(0, dims[i], factor[i]) for i in range(ndims)]
         newdims = tuple([roundup(dims[i], factor[i]) for i in range(ndims)])
 
-        return self.map(lambda v: v[slices], dims=newdims)
+        return self.map(lambda v: v[slices], value_shape=newdims)
 
     def gaussian_filter(self, sigma=2, order=0):
         """
@@ -342,7 +322,7 @@ def gaussian_filter(self, sigma=2, order=0):
         """
         from scipy.ndimage.filters import gaussian_filter
 
-        return self.map(lambda v: gaussian_filter(v, sigma, order), dims=self.dims)
+        return self.map(lambda v: gaussian_filter(v, sigma, order), value_shape=self.dims)
 
     def uniform_filter(self, size=2):
         """
@@ -415,7 +395,7 @@ def filter_(im):
         else:
             filter_ = lambda x: func(x, size)
 
-        return self.map(lambda v: filter_(v), dims=self.dims)
+        return self.map(lambda v: filter_(v), value_shape=self.dims)
 
     def localcorr(self, size=2):
         """
@@ -471,7 +451,7 @@ def subtract(self, val):
                 raise Exception('Cannot subtract image with dimensions %s '
                                 'from images with dimension %s' % (str(val.shape), str(self.dims)))
 
-        return self.map(lambda x: x - val, dims=self.dims)
+        return self.map(lambda x: x - val, value_shape=self.dims)
 
     def topng(self, path, prefix='image', overwrite=False):
         """
@@ -578,4 +558,4 @@ def map_as_series(self, func, value_size=None, dtype=None, chunk_size='auto'):
         def f(block):
             return apply_along_axis(func, 0, block)
 
-        return blocks.map(f, dims=dims, dtype=dtype).toimages()
+        return blocks.map(f, value_shape=dims, dtype=dtype).toimages()

thunder/series/series.py

@@ -170,7 +170,7 @@ def sample(self, n=100, seed=None):
 
         return self._constructor(result, index=self.index)
 
-    def map(self, func, index=None, dtype=None, with_keys=False):
+    def map(self, func, index=None, value_shape=None, dtype=None, with_keys=False):
         """
         Map an array -> array function over each record.
 
@@ -182,15 +182,33 @@ def map(self, func, index=None, dtype=None, with_keys=False):
         index : array-like, optional, default = None
             If known, the index to be used following function evaluation.
 
+        value_shape : int, optional, default=None
+            Known shape of values resulting from operation. Only
+            valid in spark mode.
+
         dtype : numpy.dtype, optional, default = None
             If known, the type of the data following function evaluation.
 
         with_keys : boolean, optional, default = False
             If true, function should be of both tuple indices and series values.
         """
-        value_shape = len(index) if index is not None else None
-        new = self._map(func, axis=self.baseaxes, value_shape=value_shape, dtype=dtype, with_keys=with_keys)
-        return self._constructor(new.values, index=index, labels=self.labels)
+        # if new index is given, can infer missing value_shape
+        if value_shape is None and index is not None:
+            value_shape = len(index)
+
+        if isinstance(value_shape, int):
+            values_shape = (value_shape, )
+        new = super(Series, self).map(func, value_shape=value_shape, dtype=dtype, with_keys=with_keys)
+
+        print(new)
+        if index is not None:
+            new.index = index
+        # if series shape did not change and no index was supplied, propagate original index
+        else:
+            if len(new.index) == len(self.index):
+                new.index = self.index
+
+        return new
 
     def reduce(self, func):
         """