Commits

Brian Kearns committed 31c7d69

fix argsort for nan/inf

  • Participants
  • Parent commits 8596a76

Comments (0)

Files changed (2)

pypy/module/micronumpy/arrayimpl/sort.py

     if count < 2:
         def arg_lt(a, b):
             # Does numpy do <= ?
-            return a[0] < b[0]
+            return a[0] < b[0] or b[0] != b[0] and a[0] == a[0]
     else:
         def arg_lt(a, b):
             for i in range(count):
+                if b[0][i] != b[0][i] and a[0][i] == a[0][i]:
+                    return True
+                elif b[0][i] == b[0][i] and a[0][i] != a[0][i]:
+                    return False
+            for i in range(count):
                 if a[0][i] < b[0][i]:
                     return True
                 elif a[0][i] > b[0][i]:

pypy/module/micronumpy/test/test_sorting.py

         assert [isnan(bb) for bb in b] == [isnan(aa) for aa in a[::-1]]
         assert (b[:2] == a[::-1][:2]).all()
 
+        b = a.argsort()
+        assert (b == [2, 1, 0]).all()
+
         # check complex
         a = zeros(9, dtype=complex128)
         a.real += [nan, nan, nan, 1, 0, 1, 1, 0, 0]
         assert [isnan(bb) for bb in b] == [isnan(aa) for aa in a[::-1]]
         assert (b[:4] == a[::-1][:4]).all()
 
+        b = a.argsort()
+        assert (b == [8, 7, 6, 5, 4, 3, 2, 1, 0]).all()
+
         # all c scalar sorts use the same code with different types
         # so it suffices to run a quick check with one type. The number
         # of sorted items must be greater than ~50 to check the actual