ENH: cumsum/cumprod refactor, consistency

Author: wesm

pandas/core/frame.py

@@ -1959,31 +1959,6 @@ def sum(self, axis=0, numeric_only=False):
 
         return Series(the_sum, index=axis_labels)
 
-    def cumsum(self, axis=0):
-        """
-        Return DataFrame of cumulative sums over requested axis.
-
-        Parameters
-        ----------
-        axis : {0, 1}
-            0 for row-wise, 1 for column-wise
-
-        Returns
-        -------
-        y : DataFrame
-        """
-        y = np.array(self.values, subok=True)
-        if not issubclass(y.dtype.type, np.int_):
-            mask = np.isnan(self.values)
-            y[mask] = 0
-            result = y.cumsum(axis)
-            has_obs = (-mask).astype(int).cumsum(axis) > 0
-            result[-has_obs] = np.nan
-        else:
-            result = y.cumsum(axis)
-        return type(self)(result, index=self.index, columns=self.columns,
-                          copy=False)
-
     def min(self, axis=0):
         """
         Return array or Series of minimums over requested axis.
@@ -2018,30 +1993,6 @@ def max(self, axis=0):
         np.putmask(values, -np.isfinite(values), -np.inf)
         return Series(values.max(axis), index=self._get_agg_axis(axis))
 
-    def cumprod(self, axis=0):
-        """
-        Return cumulative product over requested axis as DataFrame
-
-        Parameters
-        ----------
-        axis : {0, 1}
-            0 for row-wise, 1 for column-wise
-
-        Returns
-        -------
-        y : DataFrame
-        """
-        def get_cumprod(y):
-            y = np.array(y)
-            mask = isnull(y)
-            if not issubclass(y.dtype.type, np.int_):
-                y[mask] = 1
-            result = y.cumprod()
-
-            return result
-
-        return self.apply(get_cumprod, axis=axis)
-
     def product(self, axis=0):
         """
         Return array or Series of products over requested axis.

pandas/core/generic.py

@@ -166,3 +166,52 @@ def _reindex_axis(self, new_index, fill_method, axis):
             new_data = self._data.reindex_axis(new_index, axis=axis,
                                                method=fill_method)
         return type(self)(new_data)
+
+    def cumsum(self, axis=0):
+        """
+        Return DataFrame of cumulative sums over requested axis.
+
+        Parameters
+        ----------
+        axis : {0, 1}
+            0 for row-wise, 1 for column-wise
+
+        Returns
+        -------
+        y : DataFrame
+        """
+        y = self.values.copy()
+        if not issubclass(y.dtype.type, np.int_):
+            mask = np.isnan(self.values)
+            np.putmask(y, mask, 0.)
+            result = y.cumsum(axis)
+            np.putmask(result, mask, np.nan)
+        else:
+            result = y.cumsum(axis)
+        return type(self)(result, index=self.index, columns=self.columns,
+                          copy=False)
+
+    def cumprod(self, axis=0):
+        """
+        Return cumulative product over requested axis as DataFrame
+
+        Parameters
+        ----------
+        axis : {0, 1}
+            0 for row-wise, 1 for column-wise
+
+        Returns
+        -------
+        y : DataFrame
+        """
+        y = self.values.copy()
+        if not issubclass(y.dtype.type, np.int_):
+            mask = np.isnan(self.values)
+            np.putmask(y, mask, 1.)
+            result = y.cumprod(axis)
+            np.putmask(result, mask, np.nan)
+        else:
+            result = y.cumprod(axis)
+        return type(self)(result, index=self.index, columns=self.columns,
+                          copy=False)
+

pandas/core/series.py

@@ -577,21 +577,37 @@ def cumsum(self, axis=0, dtype=None, out=None):
         -------
 
         """
-        arr = self.copy()
-        okLocs = notnull(arr)
-        result = np.cumsum(arr.view(ndarray)[okLocs])
-        arr = arr.astype(result.dtype)
-        arr[okLocs] = result
-        return arr
+        arr = self.values.copy()
+
+        do_mask = not issubclass(self.dtype.type, np.int_)
+        if do_mask:
+            mask = isnull(arr)
+            np.putmask(arr, mask, 0.)
+
+        result = arr.cumsum()
+
+        if do_mask:
+            np.putmask(result, mask, np.nan)
+
+        return Series(result, index=self.index)
 
     def cumprod(self, axis=0, dtype=None, out=None):
         """
         Overriding numpy's built-in cumprod functionality
         """
-        arr = self.copy()
-        okLocs = notnull(arr)
-        arr[okLocs] = np.cumprod(arr.view(ndarray)[okLocs])
-        return arr
+        arr = self.values.copy()
+
+        do_mask = not issubclass(self.dtype.type, np.int_)
+        if do_mask:
+            mask = isnull(arr)
+            np.putmask(arr, mask, 1.)
+
+        result = arr.cumprod()
+
+        if do_mask:
+            np.putmask(result, mask, np.nan)
+
+        return Series(result, index=self.index)
 
     def median(self):
         """

pandas/core/sparse.py

@@ -10,7 +10,7 @@
 
 import operator
 
-from pandas.core.common import (isnull, _pickle_array, _unpickle_array,
+from pandas.core.common import (isnull, notnull, _pickle_array, _unpickle_array,
                                 _mut_exclusive, _ensure_index, _try_sort)
 from pandas.core.index import Index, NULL_INDEX
 from pandas.core.series import Series, TimeSeries
@@ -510,6 +510,23 @@ def sum(self, axis=None, dtype=None, out=None):
             nsparse = self.sp_index.npoints
             return sp_sum + self.fill_value * nsparse
 
+    def cumsum(self, axis=0, dtype=None, out=None):
+        """
+        Cumulative sum of values. Preserves NaN values
+
+        Extra parameters are to preserve ndarray interface.
+
+        Returns
+        -------
+
+        """
+        if not np.isnan(self.fill_value):
+            return self.to_dense().cumsum()
+        return SparseSeries(self.sp_values.cumsum(),
+                            index=self.index,
+                            sparse_index=self.sp_index,
+                            fill_value=self.fill_value)
+
     def mean(self, axis=None, dtype=None, out=None):
         """
         Mean of non-null values
@@ -1057,6 +1074,21 @@ def count(self, axis=0, **kwds):
         """
         return self.apply(SparseSeries.count, axis=axis)
 
+    def cumsum(self, axis=0):
+        """
+        Return SparseDataFrame of cumulative sums over requested axis.
+
+        Parameters
+        ----------
+        axis : {0, 1}
+            0 for row-wise, 1 for column-wise
+
+        Returns
+        -------
+        y : SparseDataFrame
+        """
+        return self.apply(SparseSeries.cumsum, axis=axis)
+
     def shift(self, periods, offset=None, timeRule=None):
         """
         Analogous to DataFrame.shift

pandas/tests/test_frame.py

@@ -2269,22 +2269,43 @@ def test_quantile(self):
         self.assertEqual(q['A'], scoreatpercentile(self.intframe['A'], 10))
 
     def test_cumsum(self):
+        self.tsframe.ix[5:10, 0] = nan
+        self.tsframe.ix[10:15, 1] = nan
+        self.tsframe.ix[15:, 2] = nan
+
+        # axis = 0
         cumsum = self.tsframe.cumsum()
-        assert_series_equal(cumsum['A'], np.cumsum(self.tsframe['A'].fillna(0)))
-        df = self.klass({'A' : np.arange(20)}, index=np.arange(20))
+        expected = self.tsframe.apply(Series.cumsum)
+        assert_frame_equal(cumsum, expected)
+
+        # axis = 1
+        cumsum = self.tsframe.cumsum(axis=1)
+        expected = self.tsframe.apply(Series.cumsum, axis=1)
+        assert_frame_equal(cumsum, expected)
 
         # works
+        df = self.klass({'A' : np.arange(20)}, index=np.arange(20))
         result = df.cumsum()
 
         # fix issue
         cumsum_xs = self.tsframe.cumsum(axis=1)
         self.assertEqual(np.shape(cumsum_xs), np.shape(self.tsframe))
 
+
     def test_cumprod(self):
+        self.tsframe.ix[5:10, 0] = nan
+        self.tsframe.ix[10:15, 1] = nan
+        self.tsframe.ix[15:, 2] = nan
+
+        # axis = 0
         cumprod = self.tsframe.cumprod()
+        expected = self.tsframe.apply(Series.cumprod)
+        assert_frame_equal(cumprod, expected)
 
-        assert_series_equal(cumprod['A'],
-                            np.cumprod(self.tsframe['A'].fillna(1)))
+        # axis = 1
+        cumprod = self.tsframe.cumprod(axis=1)
+        expected = self.tsframe.apply(Series.cumprod, axis=1)
+        assert_frame_equal(cumprod, expected)
 
         # fix issue
         cumprod_xs = self.tsframe.cumprod(axis=1)

pandas/tests/test_sparse.py

@@ -575,6 +575,11 @@ def test_shift(self):
         f = lambda s: s.shift(2, offset=datetools.bday)
         _dense_series_compare(series, f)
 
+    def test_cumsum(self):
+        result = self.bseries.cumsum()
+        expected = self.bseries.to_dense().cumsum()
+        self.assert_(isinstance(result, SparseSeries))
+        assert_series_equal(result.to_dense(), expected)
 
 class TestSparseTimeSeries(TestCase):
     pass
@@ -1066,6 +1071,12 @@ def test_count(self):
         dense_result = self.frame.to_dense().count(1)
         assert_series_equal(result, dense_result)
 
+    def test_cumsum(self):
+        result = self.frame.cumsum()
+        expected = self.frame.to_dense().cumsum()
+        self.assert_(isinstance(result, SparseDataFrame))
+        assert_frame_equal(result.to_dense(), expected)
+
     def _check_all(self, check_func):
         check_func(self.frame)
         check_func(self.iframe)