# HG changeset patch # User Justin Peel # Date 1310969792 21600 # Branch numpy-sort # Node ID 6fcaedcfdebb5d363a61d63e9d9d9885e0ebcdd6 # Parent 983ca5a89054bf435a3364d6877fc0672d092dc6 First rendition of sort. Needs more work and tests. diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py --- a/pypy/module/micronumpy/interp_numarray.py +++ b/pypy/module/micronumpy/interp_numarray.py @@ -7,6 +7,8 @@ from pypy.tool.sourcetools import func_with_new_name import math +INSERT_SORT_THRESH = 15 + def dummy1(v): assert isinstance(v, float) return v @@ -235,6 +237,80 @@ else: return self.descr_mul(space, w_other) + def _insertion_sort(self, storage, left, right): + i = left + 1 + while i <= right: + temp = storage[i] + j = i - 1 + while j >= left and storage[j] > temp: + storage[j + 1] = storage[j] + j -= 1 + storage[j + 1] = temp + i += 1 + + def descr_sort(self, space): + storage = self.get_concrete().storage + # can replace these with integer/bool numpy arrays when we add dtypes + lefts = [0] + rights = [self.find_size() - 1] + checkpivots = [False] + while lefts: + left = lefts.pop() + right = rights.pop() + checkpivot = checkpivots.pop() + # just use middle element for now. will change to med of 3 soon + mid = left + (right - left) / 2 + pivot = storage[mid] + if checkpivot and pivot == storage[left - 1]: + storage[mid], storage[left] = storage[left], storage[mid] + i = left + 1 + j = right + while 1: + while storage[j] != pivot: + j -= 1 + while storage[i] == pivot: + if i >= j: break + i += 1 + if i >= j: break + storage[i], storage[j] = storage[j], storage[i] + storage[j] = pivot + if right > j + 1: + if right - j + 1 < INSERT_SORT_THRESH: + self._insertion_sort(storage, j + 1, right) + else: + lefts.append(j + 1) + rights.append(right) + checkpivots.append(False) + else: + storage[mid], storage[right] = storage[right], storage[mid] + i = left + j = right - 1 + while 1: + while storage[i] < pivot: + i += 1 + while storage[j] >= pivot: + if i >= j: break + j -= 1 + if i >= j: break + storage[i], storage[j] = storage[j], storage[i] + storage[right] = storage[i] + storage[i] = pivot + # we can have the smaller subarray sorted first + if left < i - 1: + if i - 1 - left < INSERT_SORT_THRESH: + self._insertion_sort(storage, left, i - 1) + else: + lefts.append(left) + rights.append(i - 1) + checkpivots.append(checkpivot) + if right > i + 1: + if right - i - 1 < INSERT_SORT_THRESH: + self._insertion_sort(storage, i + 1, right) + else: + lefts.append(i + 1) + rights.append(right) + checkpivots.append(True) + def get_concrete(self): raise NotImplementedError @@ -586,4 +662,5 @@ all = interp2app(BaseArray.descr_all), any = interp2app(BaseArray.descr_any), dot = interp2app(BaseArray.descr_dot), + sort = interp2app(BaseArray.descr_sort), ) diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py --- a/pypy/module/micronumpy/test/test_numarray.py +++ b/pypy/module/micronumpy/test/test_numarray.py @@ -398,6 +398,15 @@ for i in xrange(5): assert b[i] == 2.5*a[i] + def test_sort(self): + # needs a lot more tests of sort + from numpy import array + a = array(range(19,-1,-1)) + b = array(range(20)) + a.sort() + for i in xrange(20): + assert a[i] == b[i] + class AppTestSupport(object): def setup_class(cls):