import numpy as np import numpy.random as rnd g = rnd.default_rng(1234) def less(a: T, b: T, T: type): if T is float or T is float32: return a < b or (b != b and a == a) elif T is complex or T is complex64: if a.real < b.real: return a.imag == a.imag or b.imag != b.imag elif a.real > b.real: return b.imag != b.imag and a.imag == a.imag elif a.real == b.real or (a.real != a.real and b.real != b.real): return a.imag < b.imag or (b.imag != b.imag and a.imag == a.imag) else: return b.real != b.real else: return a < b def less_equal(a, b): return a == b or less(a, b) def check_sorted(vec): for i in range(1, len(vec)): if not less_equal(vec[i - 1], vec[i]): return False return True def check_argsorted(vec, tosort): for i in range(1, len(vec)): if not less_equal(vec[tosort[i - 1]], vec[tosort[i]]): return False return True def gen_data(generator, length: int, dtype: type): if dtype is float or dtype is float32: data = generator.normal(100, 10, size=length).astype(dtype) elif dtype is complex or dtype is complex64: F = type(dtype().real) data = generator.normal(1000, 100, size=length).astype( F) + generator.normal(100, 10, size=length).astype(F) * 1j elif isinstance(dtype, UInt): data = generator.integers(0, 100000, size=length).astype(dtype) else: data = generator.integers(-100000, 100000, size=length).astype(dtype) return data @test def test_sort(generator, length: int, kind: str, dtype: type = int): print(f'TEST kind={kind} length={length} dtype={dtype.__name__}') data = gen_data(generator, length, dtype) data_copy = data.copy() assert check_sorted(np.sort(data, kind=kind)) assert check_argsorted(data, np.argsort(data, kind=kind)) assert (data == data_copy).all() data.sort(kind=kind) assert check_sorted(data) def test_sorts(dtype: type): for kind in ('quick', 'merge', 'heap'): for length in (0, 1, 10, 101, 1111, 12345, 1000000): test_sort(g, length, kind, dtype) def check_partitioned(vec, kth): if isinstance(kth, int): return check_partitioned(vec, (kth, )) s = np.sort(vec) for k in kth: if vec[k] != s[k]: return False return True def check_apartitioned(vec, kth, apart): if isinstance(kth, int): return check_apartitioned(vec, (kth, ), apart) s = np.sort(vec) for k in kth: if vec[apart[k]] != s[k]: return False return True @test def test_partition_specific(generator, length, kth): data = gen_data(generator, length, float) part = np.partition(data, kth) assert check_partitioned(part, kth) data.partition(kth) assert check_partitioned(data, kth) def test_partition(): for length in (101, 1111, 12345, 1000000): for kth in (0, 42, -1, (0, -1), (99, -99), (10, 20, 30 - 10)): test_partition_specific(g, length, kth) @test def test_apartition_specific(generator, length, kth): data = gen_data(generator, length, float) apart = np.argpartition(data, kth) assert check_apartitioned(data, kth, apart) apart = data.argpartition(kth) assert check_apartitioned(data, kth, apart) def test_apartition(): for length in (101, 1111, 12345, 1000000): for kth in (0, 42, -1, (0, -1), (99, -99), (10, 20, 30 - 10)): test_apartition_specific(g, length, kth) @test def test_sorts_along_axis(): a = np.array([[2, 22, 222], [1, 11, 111], [3, 33, 333]]) assert (np.sort(a, axis=0) == np.array([[1, 11, 111], [2, 22, 222], [3, 33, 333]])).all() assert (np.partition(a, kth=1, axis=0)[1] == np.array([2, 22, 222])).all() assert (np.argsort(a, axis=0) == np.array([[1, 1, 1], [0, 0, 0], [2, 2, 2]])).all() assert (np.argpartition(a, kth=1, axis=0)[1] == np.array([0, 0, 0])).all() assert (a.argsort(axis=0) == np.array([[1, 1, 1], [0, 0, 0], [2, 2, 2]])).all() assert (a.argpartition(kth=1, axis=0)[1] == np.array([0, 0, 0])).all() a.sort(axis=0) assert (a == np.array([[1, 11, 111], [2, 22, 222], [3, 33, 333]])).all() a = np.array([[2, 22, 222], [1, 11, 111], [3, 33, 333]]) a.partition(kth=1, axis=0) assert (a[1] == np.array([2, 22, 222])).all() @test def test_sort_complex(): assert (np.sort_complex([5, 3, 6, 2, 1]) == np.array( [1. + 0.j, 2. + 0.j, 3. + 0.j, 5. + 0.j, 6. + 0.j])).all() assert (np.sort_complex([ 1 + 2j, 2 - 1j, 3 - 2j, 3 - 3j, 3 + 5j ]) == np.array([1. + 2.j, 2. - 1.j, 3. - 3.j, 3. - 2.j, 3. + 5.j])).all() @test def test_lexsort(): a = [1, 5, 1, 4, 3, 4, 4] b = [9, 4, 0, 4, 0, 2, 1] assert np.lexsort(a) == 0 assert (np.lexsort((a, )) == np.array([0, 2, 4, 3, 5, 6, 1])).all() assert (np.lexsort((a, b)) == np.array([2, 4, 6, 5, 3, 1, 0])).all() assert (np.lexsort((b, a)) == np.array([2, 0, 4, 6, 5, 3, 1])).all() assert np.lexsort(np.array(a)) == 0 assert (np.lexsort(np.array([a])) == np.array([0, 2, 4, 3, 5, 6, 1])).all() assert (np.lexsort(np.array([a, b])) == np.array([2, 4, 6, 5, 3, 1, 0])).all() assert (np.lexsort(np.array([b, a])) == np.array([2, 0, 4, 6, 5, 3, 1])).all() assert (np.lexsort(np.array([[a, b], [a, b]])) == np.array([[0, 2, 4, 3, 5, 6, 1], [2, 4, 6, 5, 1, 3, 0]])).all() assert (np.lexsort(np.array([[a, b], [a, b]]), axis=0) == np.array([[0, 1, 1, 0, 1, 1, 1], [1, 0, 0, 1, 0, 0, 0]])).all() assert (np.lexsort(np.array([[a, b], [a, b]]).T) == np.array([[0, 1], [0, 1]])).all() assert (np.lexsort(np.array([[a, b], [a, b]]).T, axis=0) == np.array([[1, 1], [0, 0]])).all() test_sorts(int) test_sorts(u64) test_sorts(u32) test_sorts(i32) test_sorts(u8) test_sorts(i8) test_sorts(u128) test_sorts(i128) test_sorts(float) test_sorts(float32) test_sorts(complex) test_sorts(complex64) test_partition() test_apartition() test_sorts_along_axis() test_sort_complex() test_lexsort()