ENH: allow multiple axes to be passed to axis on a Panel to process slabs

Author: jreback

doc/source/release.rst

@@ -75,6 +75,7 @@ Improvements to existing features
   - Series.str.contains now has a `regex=False` keyword which can be faster for plain (non-regex) string patterns. (:issue: `5879`)
   - support ``dtypes`` on ``Panel``
   - extend ``Panel.apply`` to allow arbitrary functions (rather than only ufuncs) (:issue:`1148`)
+    allow multiple axes to be used to operate on slabs of a ``Panel``
 
 .. _release.bug_fixes-0.13.1:
 

pandas/core/panel.py

@@ -841,7 +841,7 @@ def to_frame(self, filter_observations=True):
     to_long = deprecate('to_long', to_frame)
     toLong = deprecate('toLong', to_frame)
 
-    def apply(self, func, axis='major', args=(), **kwargs):
+    def apply(self, func, axis='major', **kwargs):
         """
         Applies function along input axis of the Panel
 
@@ -852,9 +852,6 @@ def apply(self, func, axis='major', args=(), **kwargs):
             e.g. if axis = 'items', then the combination of major_axis/minor_axis
             will be passed a Series
         axis : {'major', 'minor', 'items'}
-        args : tuple
-            Positional arguments to pass to function in addition to the
-            array/series
         Additional keyword arguments will be passed as keywords to the function
 
         Examples
@@ -868,25 +865,36 @@ def apply(self, func, axis='major', args=(), **kwargs):
         -------
         result : Pandas Object
         """
-        axis = self._get_axis_number(axis)
-        axis_name = self._get_axis_name(axis)
-        ax = self._get_axis(axis)
-        values = self.values
-        ndim = self.ndim
 
-        if args or kwargs and not isinstance(func, np.ufunc):
-            f = lambda x: func(x, *args, **kwargs)
+        if kwargs and not isinstance(func, np.ufunc):
+            f = lambda x: func(x, **kwargs)
         else:
             f = func
 
+        # 2d-slabs
+        if isinstance(axis, (tuple,list)) and len(axis) == 2:
+            return self._apply_2d(f, axis=axis)
+
+        axis = self._get_axis_number(axis)
+
         # try ufunc like
         if isinstance(f, np.ufunc):
             try:
-                result = np.apply_along_axis(func, axis, values)
+                result = np.apply_along_axis(func, axis, self.values)
                 return self._wrap_result(result, axis=axis)
             except (AttributeError):
                 pass
 
+        # 1d
+        return self._apply_1d(f, axis=axis)
+
+    def _apply_1d(self, func, axis):
+
+        axis_name = self._get_axis_name(axis)
+        ax = self._get_axis(axis)
+        ndim = self.ndim
+        values = self.values
+
         # iter thru the axes
         slice_axis = self._get_axis(axis)
         slice_indexer = [0]*(ndim-1)
@@ -902,14 +910,14 @@ def apply(self, func, axis='major', args=(), **kwargs):
         points = cartesian_product(planes)
 
         results = []
-        for i in xrange(np.prod(shape)):
+        for i in range(np.prod(shape)):
 
             # construct the object
             pts = tuple([ p[i] for p in points ])
             indexer.put(indlist, slice_indexer)
 
             obj = Series(values[tuple(indexer)],index=slice_axis,name=pts)
-            result = func(obj, *args, **kwargs)
+            result = func(obj)
 
             results.append(result)
 
@@ -940,6 +948,32 @@ def apply(self, func, axis='major', args=(), **kwargs):
             planes = planes[::-1]
         return self._construct_return_type(results,planes)
 
+    def _apply_2d(self, func, axis):
+        """ handle 2-d slices, equiv to iterating over the other axis """
+
+        ndim = self.ndim
+        axis = [ self._get_axis_number(a) for a in axis ]
+
+        # construct slabs, in 2-d this is a DataFrame result
+        indexer_axis = list(range(ndim))
+        for a in axis:
+            indexer_axis.remove(a)
+        indexer_axis = indexer_axis[0]
+
+        slicer = [ slice(None,None) ] * ndim
+        ax = self._get_axis(indexer_axis)
+
+        results = []
+        for i, e in enumerate(ax):
+
+            slicer[indexer_axis] = i
+            sliced = self.iloc[tuple(slicer)]
+
+            obj = func(sliced)
+            results.append((e,obj))
+
+        return self._construct_return_type(dict(results))
+
     def _reduce(self, op, axis=0, skipna=True, numeric_only=None,
                 filter_type=None, **kwds):
         axis_name = self._get_axis_name(axis)
@@ -961,7 +995,7 @@ def _construct_return_type(self, result, axes=None, **kwargs):
         # need to assume they are the same
         if ndim is None:
             if isinstance(result,dict):
-                ndim = getattr(result.values()[0],'ndim',None)
+                ndim = getattr(list(compat.itervalues(result))[0],'ndim',None)
 
                 # a saclar result
                 if ndim is None:

pandas/tests/test_panel.py

@@ -1110,15 +1110,56 @@ def test_apply(self):
         result = self.panel.apply(lambda x: x.sum(), axis='minor_axis')
         assert_frame_equal(result,expected)
 
-        # pass args
-        result = self.panel.apply(lambda x, y: x.sum() + y, axis='items', args=[5])
+        # pass kwargs
+        result = self.panel.apply(lambda x, y: x.sum() + y, axis='items', y=5)
         expected = self.panel.sum(0) + 5
         assert_frame_equal(result,expected)
 
-        result = self.panel.apply(lambda x, y: x.sum() + y, axis='items', y=5)
-        expected = self.panel.sum(0) + 5
+    def test_apply_slabs(self):
+
+        # same shape as original
+        result = self.panel.apply(lambda x: x*2, axis = ['items','major_axis'])
+        expected = (self.panel*2).transpose('minor_axis','major_axis','items')
+        assert_panel_equal(result,expected)
+        result = self.panel.apply(lambda x: x*2, axis = ['major_axis','items'])
+        assert_panel_equal(result,expected)
+
+        result = self.panel.apply(lambda x: x*2, axis = ['items','minor_axis'])
+        expected = (self.panel*2).transpose('major_axis','minor_axis','items')
+        assert_panel_equal(result,expected)
+        result = self.panel.apply(lambda x: x*2, axis = ['minor_axis','items'])
+        assert_panel_equal(result,expected)
+
+        result = self.panel.apply(lambda x: x*2, axis = ['major_axis','minor_axis'])
+        expected = self.panel*2
+        assert_panel_equal(result,expected)
+        result = self.panel.apply(lambda x: x*2, axis = ['minor_axis','major_axis'])
+        assert_panel_equal(result,expected)
+
+        # reductions
+        result = self.panel.apply(lambda x: x.sum(0), axis = ['items','major_axis'])
+        expected = self.panel.sum(1).T
         assert_frame_equal(result,expected)
 
+        result = self.panel.apply(lambda x: x.sum(1), axis = ['items','major_axis'])
+        expected = self.panel.sum(0)
+        assert_frame_equal(result,expected)
+
+        # transforms
+        f = lambda x: (x-x.mean(1)/x.std(1))
+
+        result = self.panel.apply(f, axis = ['items','major_axis'])
+        expected = Panel(dict([ (ax,f(self.panel.loc[:,:,ax])) for ax in self.panel.minor_axis ]))
+        assert_panel_equal(result,expected)
+
+        result = self.panel.apply(f, axis = ['major_axis','minor_axis'])
+        expected = Panel(dict([ (ax,f(self.panel.loc[ax])) for ax in self.panel.items ]))
+        assert_panel_equal(result,expected)
+
+        result = self.panel.apply(f, axis = ['minor_axis','items'])
+        expected = Panel(dict([ (ax,f(self.panel.loc[:,ax])) for ax in self.panel.major_axis ]))
+        assert_panel_equal(result,expected)
+
     def test_reindex(self):
         ref = self.panel['ItemB']