fix groupby

Author: jreback

pandas/compat/numpy/function.py

@@ -306,13 +306,16 @@ def validate_expanding_func(name, args, kwargs):
             raise UnsupportedFunctionCall(msg)
 
 
-def validate_groupby_func(name, args, kwargs):
+def validate_groupby_func(name, args, kwargs, allowed_kwargs=None):
     """
-    'args' and 'kwargs' should be empty because all of
+    'args' should be empty because all of
     their necessary parameters are explicitly listed in
     the function signature
     """
-    if len(args) + len(kwargs) > 0:
+    if allowed_kwargs:
+        kwargs = set(kwargs) - set(allowed_kwargs)
+
+    if len(args) or len(kwargs):
         raise UnsupportedFunctionCall((
             "numpy operations are not valid "
             "with groupby. Use .groupby(...)."

pandas/core/groupby.py

@@ -49,6 +49,7 @@
 from pandas.formats.printing import pprint_thing
 from pandas.util.validators import validate_kwargs
 
+from pandas.tools import weightby
 import pandas.core.algorithms as algos
 import pandas.core.common as com
 from pandas.core.config import option_context
@@ -791,15 +792,21 @@ def _cython_transform(self, how, numeric_only=True):
 
         return self._wrap_transformed_output(output, names)
 
-    def _cython_agg_general(self, how, numeric_only=True):
+    def _cython_agg_general(self, how, weights=None, numeric_only=True):
+        if weights is not None:
+
+            # TODO, need to integrate this with the exclusions
+            _, weights = weightby.weightby(self.obj, weights=weights, axis=axis)
+
         output = {}
         for name, obj in self._iterate_slices():
             is_numeric = is_numeric_dtype(obj.dtype)
             if numeric_only and not is_numeric:
                 continue
 
+            values = weightby.weight(obj.values, weights)
             try:
-                result, names = self.grouper.aggregate(obj.values, how)
+                result, names = self.grouper.aggregate(values, how)
             except AssertionError as e:
                 raise GroupByError(str(e))
             output[name] = self._try_cast(result, obj)
@@ -1006,6 +1013,26 @@ def count(self):
         # defined here for API doc
         raise NotImplementedError
 
+    @Substitution(name='groupby')
+    @Appender(_doc_template)
+    def sum(self, *args, **kwargs):
+        """
+        Compute sum of groups, excluding missing values
+
+        For multiple groupings, the result index will be a MultiIndex
+        """
+
+        # TODO: this is slightly different from other cythonized functions (e.g. mean)
+        # to accomodate np.sum functionaility
+        nv.validate_groupby_func('sum', args, kwargs, ('weights', 'numeric_only'))
+        self._set_group_selection()
+        try:
+            return self._cython_agg_general('add', **kwargs)
+        except AssertionError as e:
+            raise SpecificationError(str(e))
+        except Exception:  # pragma: no cover
+            return self.aggregate(lambda x: np.sum(x, axis=self.axis))
+
     @Substitution(name='groupby')
     @Appender(_doc_template)
     def mean(self, *args, **kwargs):
@@ -1014,14 +1041,15 @@ def mean(self, *args, **kwargs):
 
         For multiple groupings, the result index will be a MultiIndex
         """
-        nv.validate_groupby_func('mean', args, kwargs)
+        nv.validate_groupby_func('mean', args, kwargs, ('weights', 'numeric_only'))
         try:
-            return self._cython_agg_general('mean')
+            return self._cython_agg_general('mean', **kwargs)
         except GroupByError:
             raise
         except Exception:  # pragma: no cover
             self._set_group_selection()
-            f = lambda x: x.mean(axis=self.axis)
+            kwargs['axis'] = self.axis
+            f = lambda x: x.mean(**kwargs)
             return self._python_agg_general(f)
 
     @Substitution(name='groupby')
@@ -1107,7 +1135,6 @@ def size(self):
         """Compute group sizes"""
         return self.grouper.size()
 
-    sum = _groupby_function('sum', 'add', np.sum)
     prod = _groupby_function('prod', 'prod', np.prod)
     min = _groupby_function('min', 'min', np.min, numeric_only=False)
     max = _groupby_function('max', 'max', np.max, numeric_only=False)
@@ -3134,9 +3161,9 @@ def _iterate_slices(self):
                 continue
             yield val, slicer(val)
 
-    def _cython_agg_general(self, how, numeric_only=True):
+    def _cython_agg_general(self, how, **kwargs):
         new_items, new_blocks = self._cython_agg_blocks(
-            how, numeric_only=numeric_only)
+            how, **kwargs)
         return self._wrap_agged_blocks(new_items, new_blocks)
 
     def _wrap_agged_blocks(self, items, blocks):
@@ -3162,18 +3189,26 @@ def _wrap_agged_blocks(self, items, blocks):
 
     _block_agg_axis = 0
 
-    def _cython_agg_blocks(self, how, numeric_only=True):
+    def _cython_agg_blocks(self, how, weights=None, numeric_only=True):
         data, agg_axis = self._get_data_to_aggregate()
 
+        if weights is not None:
+
+            # TODO, need to integrate this with the exclusions
+            _, weights = weightby.weightby(self.obj,
+                                           weights=weights,
+                                           axis=self.axis)
+
         new_blocks = []
 
         if numeric_only:
             data = data.get_numeric_data(copy=False)
 
         for block in data.blocks:
 
+            values = weightby.weight(block.values, weights)
             result, _ = self.grouper.aggregate(
-                block.values, how, axis=agg_axis)
+                values, how, axis=agg_axis)
 
             # see if we can cast the block back to the original dtype
             result = block._try_coerce_and_cast_result(result)

pandas/core/nanops.py

@@ -23,7 +23,7 @@
 from pandas.types.missing import isnull, notnull
 
 from pandas.core.common import _values_from_object
-
+from pandas.tools import weightby
 
 class disallow(object):
     def __init__(self, *dtypes):
@@ -204,9 +204,7 @@ def _get_values(values, skipna,
         mask = isnull(values)
 
     # weights
-    if weights is not None:
-        values = values * weights.reshape(values.shape)
-
+    values = weightby.weight(values, weights)
     dtype = values.dtype
     dtype_ok = _na_ok_dtype(dtype)
 

pandas/tools/tests/test_weightby.py

@@ -13,6 +13,9 @@ def setUp(self):
                              'B': [1, 2, 3, 4]})
         self.df2 = DataFrame({'A': [1, 2, 3, 4],
                               'B': [1, 2, 3, 4]})
+        self.df3 = DataFrame({'A': [1, 2, 3, 4],
+                              'B': [1, 2, 3, 4],
+                              'C': [1, 1, 2, 2]})
 
     @property
     def rs(self):
@@ -55,6 +58,16 @@ def test_basic(self):
             expected = getattr(self.df2[['B']] * weights2, f)(ddof=2)
             # tm.assert_series_equal(result, expected)
 
+    def test_groupby(self):
+
+        for f in ['mean', 'sum']:
+            weights = (self.df3['A'] / self.df3.A.sum()).values
+            result = getattr(self.df3.groupby('C'), f)(weights='A')
+            adj = self.df3.assign(A=self.df3.A * weights,
+                                  B=self.df3.B * weights)
+            expected = getattr(adj.groupby('C'), f)()
+            tm.assert_frame_equal(result, expected)
+
     def test_unsupported(self):
         for f in ['first', 'median', 'min', 'max', 'prod']:
 

pandas/tools/weightby.py

@@ -94,3 +94,30 @@ def weightby(obj, weights=None, axis=0):
             raise ValueError("Invalid weights: weights sum to zero")
 
     return obj, weights.values
+
+
+def weight(values, weights):
+    """
+    Return the values * weights, broadcasting if needed
+
+    Parameters
+    ----------
+    values : ndarray
+    weights : 1d-ndarray
+
+    Returns
+    -------
+    values shaped ndarray
+    """
+
+    if weights is None:
+        return values
+
+    if values.ndim == 1:
+        return values * weights
+
+    elif values.ndim == 2:
+
+        return values * weights
+
+    raise NotImplementedError