Author: Brian Kearns <[email protected]>
Branch:
Changeset: r67467:31c7d69c6043
Date: 2013-10-17 17:11 -0400
http://bitbucket.org/pypy/pypy/changeset/31c7d69c6043/
Log: fix argsort for nan/inf
diff --git a/pypy/module/micronumpy/arrayimpl/sort.py
b/pypy/module/micronumpy/arrayimpl/sort.py
--- a/pypy/module/micronumpy/arrayimpl/sort.py
+++ b/pypy/module/micronumpy/arrayimpl/sort.py
@@ -100,10 +100,15 @@
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]:
diff --git a/pypy/module/micronumpy/test/test_sorting.py
b/pypy/module/micronumpy/test/test_sorting.py
--- a/pypy/module/micronumpy/test/test_sorting.py
+++ b/pypy/module/micronumpy/test/test_sorting.py
@@ -108,6 +108,9 @@
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]
@@ -117,6 +120,9 @@
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
_______________________________________________
pypy-commit mailing list
[email protected]
https://mail.python.org/mailman/listinfo/pypy-commit
