BUG: multi-key groupby now works with multiple functions. starting to add *args, **kwargs. address GH #133

Author: wesm

pandas/core/groupby.py

@@ -218,7 +218,7 @@ def flatten(gen, level=0, shape_axis=0):
 
         return flatten(self._generator_factory(data), shape_axis=self.axis)
 
-    def apply(self, func):
+    def apply(self, func, *args, **kwargs):
         """
         Apply function and combine results together in an intelligent way. The
         split-apply-combine combination rules attempt to be as common sense
@@ -255,16 +255,16 @@ def apply(self, func):
         -------
         applied : type depending on grouped object and function
         """
-        return self._python_apply_general(func)
+        return self._python_apply_general(func, *args, **kwargs)
 
-    def aggregate(self, func):
+    def aggregate(self, func, *args, **kwargs):
         raise NotImplementedError
 
-    def agg(self, func):
+    def agg(self, func, *args, **kwargs):
         """
         See docstring for aggregate
         """
-        return self.aggregate(func)
+        return self.aggregate(func, *args, **kwargs)
 
     def _get_names(self):
         axes = [ping.group_index for ping in self.groupings]
@@ -275,18 +275,18 @@ def _get_names(self):
     def _iterate_slices(self):
         yield self.name, self.obj
 
-    def transform(self, func):
+    def transform(self, func, *args, **kwargs):
         raise NotImplementedError
 
-    def mean(self, axis=None):
+    def mean(self):
         """
         Compute mean of groups, excluding missing values
 
         For multiple groupings, the result index will be a MultiIndex
         """
         return self._cython_agg_general('mean')
 
-    def sum(self, axis=None):
+    def sum(self):
         """
         Compute sum of values, excluding missing values
 
@@ -330,7 +330,7 @@ def _get_group_levels(self, mask):
         name_list = self._get_names()
         return [(name, raveled[mask]) for name, raveled in name_list]
 
-    def _python_agg_general(self, arg):
+    def _python_agg_general(self, func, *args, **kwargs):
         group_shape = self._group_shape
         counts = np.zeros(group_shape, dtype=int)
 
@@ -342,7 +342,7 @@ def _doit(reschunk, ctchunk, gen, shape_axis=0):
                     ctchunk[i] = size
                     if size == 0:
                         continue
-                    reschunk[i] = arg(subgen)
+                    reschunk[i] = func(subgen, *args, **kwargs)
                 else:
                     _doit(reschunk[i], ctchunk[i], subgen,
                           shape_axis=shape_axis)
@@ -378,14 +378,14 @@ def _doit(reschunk, ctchunk, gen, shape_axis=0):
 
         return self._wrap_aggregated_output(output, mask)
 
-    def _python_apply_general(self, arg):
+    def _python_apply_general(self, func, *args, **kwargs):
         result_keys = []
         result_values = []
 
         not_indexed_same = False
         for key, group in self:
             group.name = key
-            res = arg(group)
+            res = func(group, *args, **kwargs)
             if not _is_indexed_like(res, group):
                 not_indexed_same = True
 
@@ -588,7 +588,7 @@ class SeriesGroupBy(GroupBy):
     def _agg_stride_shape(self):
         return ()
 
-    def aggregate(self, func_or_funcs):
+    def aggregate(self, func_or_funcs, *args, **kwargs):
         """
         Apply aggregation function or functions to groups, yielding most likely
         Series but in some cases DataFrame depending on the output of the
@@ -640,18 +640,18 @@ def aggregate(self, func_or_funcs):
         Series or DataFrame
         """
         if isinstance(func_or_funcs, basestring):
-            return getattr(self, func_or_funcs)()
-
-        if len(self.groupings) > 1:
-            return self._python_agg_general(func_or_funcs)
+            return getattr(self, func_or_funcs)(*args, **kwargs)
 
         if hasattr(func_or_funcs,'__iter__'):
             ret = self._aggregate_multiple_funcs(func_or_funcs)
         else:
+            if len(self.groupings) > 1:
+                return self._python_agg_general(func_or_funcs, *args, **kwargs)
+
             try:
-                result = self._aggregate_simple(func_or_funcs)
+                result = self._aggregate_simple(func_or_funcs, *args, **kwargs)
             except Exception:
-                result = self._aggregate_named(func_or_funcs)
+                result = self._aggregate_named(func_or_funcs, *args, **kwargs)
 
             if len(result) > 0:
                 if isinstance(result.values()[0], Series):
@@ -711,34 +711,30 @@ def _aggregate_multiple_funcs(self, arg):
         results = {}
 
         for name, func in arg.iteritems():
-            try:
-                result = func(self)
-            except Exception:
-                result = self.aggregate(func)
-            results[name] = result
+            results[name] = self.aggregate(func)
 
         return DataFrame(results)
 
-    def _aggregate_simple(self, arg):
+    def _aggregate_simple(self, func, *args, **kwargs):
         values = self.obj.values
         result = {}
         for k, v in self.primary.indices.iteritems():
-            result[k] = arg(values.take(v))
+            result[k] = func(values.take(v), *args, **kwargs)
 
         return result
 
-    def _aggregate_named(self, arg):
+    def _aggregate_named(self, func, *args, **kwargs):
         result = {}
 
         for name in self.primary:
             grp = self.get_group(name)
             grp.name = name
-            output = arg(grp)
+            output = func(grp, *args, **kwargs)
             result[name] = output
 
         return result
 
-    def transform(self, func):
+    def transform(self, func, *args, **kwargs):
         """
         Call function producing a like-indexed Series on each group and return
         a Series with the transformed values
@@ -760,7 +756,7 @@ def transform(self, func):
 
         for name, group in self:
             group.name = name
-            res = func(group)
+            res = func(group, *args, **kwargs)
             indexer, _ = self.obj.index.get_indexer(group.index)
             np.put(result, indexer, res)
 
@@ -817,7 +813,7 @@ def _obj_with_exclusions(self):
         else:
             return self.obj
 
-    def aggregate(self, arg):
+    def aggregate(self, arg, *args, **kwargs):
         """
         Aggregate using input function or dict of {column -> function}
 
@@ -843,26 +839,27 @@ def aggregate(self, arg):
             result = DataFrame(result)
         else:
             if len(self.groupings) > 1:
-                return self._python_agg_general(arg)
-            result = self._aggregate_generic(arg, axis=self.axis)
+                return self._python_agg_general(arg, *args, **kwargs)
+            result = self._aggregate_generic(arg, *args, **kwargs)
 
         return result
 
-    def _aggregate_generic(self, agger, axis=0):
+    def _aggregate_generic(self, func, *args, **kwargs):
         result = {}
-
+        axis = self.axis
         obj = self._obj_with_exclusions
 
         try:
             for name in self.primary:
                 data = self.get_group(name, obj=obj)
                 try:
-                    result[name] = agger(data)
+                    result[name] = func(data, *args, **kwargs)
                 except Exception:
-                    result[name] = data.apply(agger, axis=axis)
+                    wrapper = lambda x: func(x, *args, **kwargs)
+                    result[name] = data.apply(wrapper, axis=axis)
         except Exception, e1:
             if axis == 0:
-                return self._aggregate_item_by_item(agger)
+                return self._aggregate_item_by_item(func, *args, **kwargs)
             else:
                 raise e1
 
@@ -872,7 +869,7 @@ def _aggregate_generic(self, agger, axis=0):
 
         return result
 
-    def _aggregate_item_by_item(self, agger):
+    def _aggregate_item_by_item(self, func, *args, **kwargs):
         # only for axis==0
 
         obj = self._obj_with_exclusions
@@ -881,7 +878,7 @@ def _aggregate_item_by_item(self, agger):
         cannot_agg = []
         for item in obj:
             try:
-                result[item] = self[item].agg(agger)
+                result[item] = self[item].agg(func, *args, **kwargs)
             except (ValueError, TypeError):
                 cannot_agg.append(item)
                 continue
@@ -954,15 +951,15 @@ def _wrap_applied_output(self, keys, values, not_indexed_same=False):
 
             return DataFrame(stacked_values, index=index, columns=columns)
 
-    def transform(self, func):
+    def transform(self, func, *args, **kwargs):
         """
         Call function producing a like-indexed DataFrame on each group and
         return a DataFrame having the same indexes as the original object
         filled with the transformed values
 
         Parameters
         ----------
-        func : function
+        f : function
             Function to apply to each subframe
 
         Note
@@ -982,9 +979,10 @@ def transform(self, func):
             group.name = name
 
             try:
-                res = group.apply(func, axis=self.axis)
+                wrapper = lambda x: func(x, *args, **kwargs)
+                res = group.apply(wrapper, axis=self.axis)
             except Exception: # pragma: no cover
-                res = func(group)
+                res = func(group, *args, **kwargs)
 
             # broadcasting
             if isinstance(res, Series):
@@ -1081,7 +1079,7 @@ def _all_indexes_same(indexes):
 
 class PanelGroupBy(GroupBy):
 
-    def aggregate(self, func):
+    def aggregate(self, func, *args, **kwargs):
         """
         Aggregate using input function or dict of {column -> function}
 
@@ -1096,19 +1094,22 @@ def aggregate(self, func):
         -------
         aggregated : Panel
         """
-        return self._aggregate_generic(func, axis=self.axis)
+        return self._aggregate_generic(func, *args, **kwargs)
 
-    def _aggregate_generic(self, agger, axis=0):
+    def _aggregate_generic(self, func, *args, **kwargs):
         result = {}
 
+        axis = self.axis
+
         obj = self._obj_with_exclusions
 
         for name in self.primary:
             data = self.get_group(name, obj=obj)
             try:
-                result[name] = agger(data)
+                result[name] = func(data, *args, **kwargs)
             except Exception:
-                result[name] = data.apply(agger, axis=axis)
+                wrapper = lambda x: func(x, *args, **kwargs)
+                result[name] = data.apply(wrapper, axis=axis)
 
         result = Panel.fromDict(result, intersect=False)
 

pandas/tests/test_groupby.py

@@ -475,6 +475,14 @@ def test_multi_func(self):
         grouped = df.groupby(['k1', 'k2'])
         grouped.agg(np.sum)
 
+    def test_multi_key_multiple_functions(self):
+        grouped = self.df.groupby(['A', 'B'])['C']
+
+        agged = grouped.agg([np.mean, np.std])
+        expected = DataFrame({'mean' : grouped.agg(np.mean),
+                              'std' : grouped.agg(np.std)})
+        assert_frame_equal(agged, expected)
+
     def test_groupby_multiple_columns(self):
         data = self.df
         grouped = data.groupby(['A', 'B'])