Adding argmax and argmin with proper behavior (pandas-dev#16830)

Author: Lucas Kushner

pandas/core/frame.py

@@ -5329,7 +5329,7 @@ def idxmin(self, axis=0, skipna=True):
 
     def idxmax(self, axis=0, skipna=True):
         """
-        Return index of first occurrence of maximum over requested axis.
+        Return label of first occurrence of maximum over requested axis.
         NA/null values are excluded.
 
         Parameters
@@ -5358,6 +5358,64 @@ def idxmax(self, axis=0, skipna=True):
         result = [index[i] if i >= 0 else NA for i in indices]
         return Series(result, index=self._get_agg_axis(axis))
 
+    def argmin(self, axis=0, skipna=True):
+        """
+        Return index of first occurrence of minimum over requested axis.
+        NA/null values are excluded.
+
+        Parameters
+        ----------
+        axis : {0 or 'index', 1 or 'columns'}, default 0
+            0 or 'index' for row-wise, 1 or 'columns' for column-wise
+        skipna : boolean, default True
+            Exclude NA/null values. If an entire row/column is NA, the result
+            will be NA
+
+        Returns
+        -------
+        argmin : Series
+
+        Notes
+        -----
+        This method is the DataFrame version of ``ndarray.argmin``.
+
+        See Also
+        --------
+        Series.idxmin
+        """
+        axis = self._get_axis_number(axis)
+        indices = nanops.nanargmin(self.values, axis=axis, skipna=skipna)
+        return Series(indices, index=self._get_agg_axis(axis))
+
+    def argmax(self, axis=0, skipna=True):
+        """
+        Return index of first occurrence of maximum over requested axis.
+        NA/null values are excluded.
+
+        Parameters
+        ----------
+        axis : {0 or 'index', 1 or 'columns'}, default 0
+            0 or 'index' for row-wise, 1 or 'columns' for column-wise
+        skipna : boolean, default True
+            Exclude NA/null values. If an entire row/column is NA, the result
+            will be first index.
+
+        Returns
+        -------
+        argmax : Series
+
+        Notes
+        -----
+        This method is the DataFrame version of ``ndarray.argmax``.
+
+        See Also
+        --------
+        Series.argmax
+        """
+        axis = self._get_axis_number(axis)
+        indices = nanops.nanargmax(self.values, axis=axis, skipna=skipna)
+        return Series(indices, index=self._get_agg_axis(axis))
+
     def _get_agg_axis(self, axis_num):
         """ let's be explict about this """
         if axis_num == 0:

pandas/core/series.py

@@ -72,7 +72,7 @@
 import pandas.core.nanops as nanops
 import pandas.io.formats.format as fmt
 from pandas.util._decorators import (
-    Appender, deprecate, deprecate_kwarg, Substitution)
+    Appender, deprecate_kwarg, Substitution)
 from pandas.util._validators import validate_bool_kwarg
 
 from pandas._libs import index as libindex, tslib as libts, lib, iNaT
@@ -1239,7 +1239,7 @@ def duplicated(self, keep='first'):
 
     def idxmin(self, axis=None, skipna=True, *args, **kwargs):
         """
-        Index of first occurrence of minimum of values.
+        Label of first occurrence of minimum of values.
 
         Parameters
         ----------
@@ -1259,15 +1259,14 @@ def idxmin(self, axis=None, skipna=True, *args, **kwargs):
         DataFrame.idxmin
         numpy.ndarray.argmin
         """
-        skipna = nv.validate_argmin_with_skipna(skipna, args, kwargs)
-        i = nanops.nanargmin(_values_from_object(self), skipna=skipna)
+        i = self.argmin(axis, skipna, *args, **kwargs)
         if i == -1:
             return np.nan
         return self.index[i]
 
     def idxmax(self, axis=None, skipna=True, *args, **kwargs):
         """
-        Index of first occurrence of maximum of values.
+        Label of first occurrence of maximum of values.
 
         Parameters
         ----------
@@ -1287,15 +1286,62 @@ def idxmax(self, axis=None, skipna=True, *args, **kwargs):
         DataFrame.idxmax
         numpy.ndarray.argmax
         """
-        skipna = nv.validate_argmax_with_skipna(skipna, args, kwargs)
-        i = nanops.nanargmax(_values_from_object(self), skipna=skipna)
+        i = self.argmax(axis, skipna, *args, **kwargs)
         if i == -1:
             return np.nan
         return self.index[i]
 
-    # ndarray compat
-    argmin = deprecate('argmin', idxmin)
-    argmax = deprecate('argmax', idxmax)
+    def argmin(self, axis=None, skipna=True, *args, **kwargs):
+        """
+        Index of first occurrence of minimum of values.
+
+        Parameters
+        ----------
+        skipna : boolean, default True
+            Exclude NA/null values
+
+        Returns
+        -------
+        idxmin : Index of minimum of values
+
+        Notes
+        -----
+        This method is the Series version of ``ndarray.argmin``.
+
+        See Also
+        --------
+        DataFrame.argmin
+        numpy.ndarray.argmin
+        """
+        skipna = nv.validate_argmin_with_skipna(skipna, args, kwargs)
+        i = nanops.nanargmin(_values_from_object(self), skipna=skipna)
+        return i
+
+    def argmax(self, axis=None, skipna=True, *args, **kwargs):
+        """
+        Index of first occurrence of maximum of values.
+
+        Parameters
+        ----------
+        skipna : boolean, default True
+            Exclude NA/null values
+
+        Returns
+        -------
+        idxmax : Index of maximum of values
+
+        Notes
+        -----
+        This method is the Series version of ``ndarray.argmax``.
+
+        See Also
+        --------
+        DataFrame.argmax
+        numpy.ndarray.argmax
+        """
+        skipna = nv.validate_argmax_with_skipna(skipna, args, kwargs)
+        i = nanops.nanargmax(_values_from_object(self), skipna=skipna)
+        return i
 
     def round(self, decimals=0, *args, **kwargs):
         """

pandas/tests/series/test_analytics.py

@@ -1215,6 +1215,7 @@ def test_numpy_argmin(self):
         data = np.random.randint(0, 11, size=10)
         result = np.argmin(Series(data))
         assert result == np.argmin(data)
+        assert result == Series(data).argmin()
 
         if not _np_version_under1p10:
             msg = "the 'out' parameter is not supported"
@@ -1271,6 +1272,7 @@ def test_numpy_argmax(self):
         data = np.random.randint(0, 11, size=10)
         result = np.argmax(Series(data))
         assert result == np.argmax(data)
+        assert result == Series(data).argmax()
 
         if not _np_version_under1p10:
             msg = "the 'out' parameter is not supported"