Merge pull request #26716 from JuliaPoo/cont-15128-topk-feat

mattip · web-flow · commit c21ac104e544 · 2024-07-01T10:53:46.000+03:00
DOC: Add clarifications np.argpartition
diff --git a/numpy/_core/fromnumeric.py b/numpy/_core/fromnumeric.py
@@ -823,6 +823,8 @@ def partition(a, kth, axis=-1, kind='introselect', order=None):
     the real parts except when they are equal, in which case the order
     is determined by the imaginary parts.
 
+    The sort order of ``np.nan`` is bigger than ``np.inf``.
+
     Examples
     --------
     >>> a = np.array([7, 1, 7, 7, 1, 5, 7, 2, 3, 2, 6, 2, 3, 0])
@@ -918,7 +920,14 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
 
     Notes
     -----
-    See `partition` for notes on the different selection algorithms.
+    The returned indices are not guaranteed to be sorted according to
+    the values. Furthermore, the default selection algorithm ``introselect``
+    is unstable, and hence the returned indices are not guaranteed
+    to be the earliest/latest occurrence of the element.
+
+    `argpartition` works for real/complex inputs with nan values,
+    see `partition` for notes on the enhanced sort order and
+    different selection algorithms.
 
     Examples
     --------
diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
@@ -10211,6 +10211,16 @@ def test_partition_fp(N, dtype):
     assert_arr_partitioned(np.sort(arr)[k], k,
             arr[np.argpartition(arr, k, kind='introselect')])
 
+    # Check that `np.inf < np.nan`
+    # This follows np.sort
+    arr[0] = np.nan
+    arr[1] = np.inf
+    o1 = np.partition(arr, -2, kind='introselect')
+    o2 = arr[np.argpartition(arr, -2, kind='introselect')]
+    for out in [o1, o2]:
+        assert_(np.isnan(out[-1]))
+        assert_equal(out[-2], np.inf)
+
 def test_cannot_assign_data():
     a = np.arange(10)
     b = np.linspace(0, 1, 10)
