Add axis parameter to clip, clip_lower, and clip_upper. Fix the broken frame-clipping test.

Author: evanpw

pandas/core/generic.py

@@ -2801,14 +2801,16 @@ def notnull(self):
         """
         return notnull(self).__finalize__(self)
 
-    def clip(self, lower=None, upper=None, out=None):
+    def clip(self, lower=None, upper=None, out=None, axis=None):
         """
         Trim values at input threshold(s)
 
         Parameters
         ----------
         lower : float or array_like, default None
         upper : float or array_like, default None
+        axis : int, optional
+            Align object with lower and upper along the given axis.
 
         Returns
         -------
@@ -2824,19 +2826,21 @@ def clip(self, lower=None, upper=None, out=None):
 
         result = self
         if lower is not None:
-            result = result.clip_lower(lower)
+            result = result.clip_lower(lower, axis)
         if upper is not None:
-            result = result.clip_upper(upper)
+            result = result.clip_upper(upper, axis)
 
         return result
 
-    def clip_upper(self, threshold):
+    def clip_upper(self, threshold, axis=None):
         """
         Return copy of input with values above given value(s) truncated
 
         Parameters
         ----------
         threshold : float or array_like
+        axis : int, optional
+            Align object with threshold along the given axis.
 
         See also
         --------
@@ -2849,15 +2853,22 @@ def clip_upper(self, threshold):
         if np.any(isnull(threshold)):
             raise ValueError("Cannot use an NA value as a clip threshold")
 
-        return self.where(self.le(threshold, axis=0) | isnull(self), threshold, axis=0)
+        if axis is None:
+            return self.where(self.le(threshold) | isnull(self), threshold)
+        else:
+            subset = self.le(threshold, axis=axis) | isnull(self)
+            return self.where(subset, threshold, axis=axis)
+
 
-    def clip_lower(self, threshold):
+    def clip_lower(self, threshold, axis=None):
         """
         Return copy of the input with values below given value(s) truncated
 
         Parameters
         ----------
         threshold : float or array_like
+        axis : int, optional
+            Align object with threshold along the given axis.
 
         See also
         --------
@@ -2870,7 +2881,12 @@ def clip_lower(self, threshold):
         if np.any(isnull(threshold)):
             raise ValueError("Cannot use an NA value as a clip threshold")
 
-        return self.where(self.ge(threshold, axis=0) | isnull(self), threshold, axis=0)
+        if axis is None:
+            return self.where(self.ge(threshold) | isnull(self), threshold)
+        else:
+            subset = self.ge(threshold, axis=axis) | isnull(self)
+            return self.where(subset, threshold, axis=axis)
+
 
     def groupby(self, by=None, axis=0, level=None, as_index=True, sort=True,
                 group_keys=True, squeeze=False):

pandas/tests/test_frame.py

@@ -11258,15 +11258,12 @@ def test_clip_against_series(self):
         lb = Series(np.random.randn(1000))
         ub = lb + 1
 
-        clipped_df = df.clip(lb, ub)
+        clipped_df = df.clip(lb, ub, axis=0)
 
         for i in xrange(2):
             lb_mask = df.iloc[:, i] <= lb
             ub_mask = df.iloc[:, i] >= ub
             mask = ~lb_mask & ~ub_mask
-            #self.assertTrue((clipped_df.loc[lb_mask, i] == lb.loc[lb_mask]).all())
-            #self.assertTrue((clipped_df.loc[ub_mask, i] == ub.loc[ub_mask]).all())
-            #self.assertTrue((clipped_df.loc[mask, i] == df.loc[mask, i]).all())
 
             assert_series_equal(clipped_df.loc[lb_mask, i], lb[lb_mask])
             assert_series_equal(clipped_df.loc[ub_mask, i], ub[ub_mask])