Merge pull request #459 from nmearl/master

Let apply_function method return cube-like objects; custom update function

spectral_cube/spectral_cube.py

@@ -776,7 +776,8 @@ def _cube_on_cube_operation(self, function, cube, equivalencies=[], **kwargs):
         return self._new_cube_with(data=data, unit=unit)
 
     def apply_function(self, function, axis=None, weights=None, unit=None,
-                       projection=False, progressbar=False, **kwargs):
+                       projection=False, progressbar=False, 
+                       update_function=None, keep_shape=False, **kwargs):
         """
         Apply a function to valid data along the specified axis or to the whole
         cube, optionally using a weight array that is the same shape (or at
@@ -800,6 +801,13 @@ def apply_function(self, function, axis=None, weights=None, unit=None,
         progressbar : bool
             Show a progressbar while iterating over the slices/rays through the
             cube?
+        keep_shape : bool
+            If `True`, the returned object will be the same dimensionality as
+            the cube.
+        update_function : function
+            An alternative tracker for the progress of applying the function
+            to the cube data. If ``progressbar`` is ``True``, this argument is
+            ignored.
 
         Returns
         -------
@@ -822,13 +830,16 @@ def apply_function(self, function, axis=None, weights=None, unit=None,
 
         # determine the output array shape
         nx, ny = self._get_flat_shape(axis)
+        nz = self.shape[axis] if keep_shape else 1
 
         # allocate memory for output array
-        out = np.empty([nx, ny]) * np.nan
+        out = np.empty([nz, nx, ny]) * np.nan
 
         if progressbar:
             progressbar = ProgressBar(nx*ny)
             pbu = progressbar.update
+        elif update_function is not None:
+            pbu = update_function
         else:
             pbu = lambda: True
 
@@ -840,12 +851,15 @@ def apply_function(self, function, axis=None, weights=None, unit=None,
                 result = function(data, **kwargs)
                 if hasattr(result, 'value'):
                     # store result in array
-                    out[y, x] = result.value
+                    out[:, y, x] = result.value
                 else:
-                    out[y, x] = result
+                    out[:, y, x] = result
             pbu()
 
-        if projection and axis in (0,1,2):
+        if not keep_shape:
+            out = out[0, :, :]
+
+        if projection and axis in (0, 1, 2):
             new_wcs = wcs_utils.drop_axis(self._wcs, np2wcs[axis])
 
             meta = {'collapse_axis': axis}