Repository: numpy/numpy
File: numpy/_core/tests/test_multiarray.py (Python)

Code before:


@pytest.mark.parametrize("N", np.arange(2, 512))
@pytest.mark.parametrize("dtype", [np.int32, np.uint32, np.int64, np.uint64])
def test_argsort_int(N, dtype):
    rnd = np.random.RandomState(1100710816)
    # (1) random data with min and max values
    minv = np.iinfo(dtype).min
    maxv = np.iinfo(dtype).max
    arr = rnd.randint(low=minv, high=maxv, size=N, dtype=dtype)
    i, j = rnd.choice(N, 2, replace=False)
    arr[i] = minv
    arr[j] = maxv
    assert_arg_sorted(arr, np.argsort(arr, kind='quick'))

    # (2) random data with max value at the end of array
    # See: https://github.com/intel/x86-simd-sort/pull/39
    arr = rnd.randint(low=minv, high=maxv, size=N, dtype=dtype)
    arr[N - 1] = maxv
    assert_arg_sorted(arr, np.argsort(arr, kind='quick'))

# Test large arrays that leverage openMP implementations from x86-simd-sort:
@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
def test_sort_largearrays(dtype):
    N = 1000000
    arr = np.random.rand(N)
    assert_equal(np.sort(arr, kind='quick'), np.sort(arr, kind='heap'))

@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
def test_gh_22683():
    b = 777.68760986
    a = np.array([b] * 10000, dtype=object)
    refc_start = sys.getrefcount(b)
    np.choose(np.zeros(10000, dtype=int), [a], out=a)
    np.choose(np.zeros(10000, dtype=int), [a], out=a)
    refc_end = sys.getrefcount(b)
    assert refc_end - refc_start < 10


def test_gh_24459():
    a = np.zeros((50, 3), dtype=np.float64)
    with pytest.raises(TypeError):
        np.choose(a, [3, -1])


def test_gh_28206():
    a = np.arange(3)
    b = np.ones((3, 3), dtype=np.int64)
    out = np.array([np.nan, np.nan, np.nan])

    with warnings.catch_warnings():

Review comment (performance):
maybe we should pin the values down with modern `default_rng`?

One other thing I checked was how slow the test might be (do we want a `slow` marker for the large array handling?). It didn't seem too bad (< 1 s for each case locally on an ARM Mac), though it is in the top 10 slowest for this module for example.

Code after:
    arr = rnd.randint(low=minv, high=maxv, size=N, dtype=dtype)
    i, j = rnd.choice(N, 2, replace=False)
    arr[i] = minv
    arr[j] = maxv
    assert_arg_sorted(arr, np.argsort(arr, kind='quick'))

    # (2) random data with max value at the end of array
    # See: https://github.com/intel/x86-simd-sort/pull/39
    arr = rnd.randint(low=minv, high=maxv, size=N, dtype=dtype)
    arr[N - 1] = maxv
    assert_arg_sorted(arr, np.argsort(arr, kind='quick'))

# Test large arrays that leverage openMP implementations from x86-simd-sort:
@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
def test_sort_largearrays(dtype):
    N = 1000000
    rnd = np.random.RandomState(1100710816)
    arr = -0.5 + rnd.random(N).astype(dtype)
    assert_equal(np.sort(arr, kind='quick'), np.sort(arr, kind='heap'))

# Test large arrays that leverage openMP implementations from x86-simd-sort:
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
def test_argsort_largearrays(dtype):
    N = 1000000
    rnd = np.random.RandomState(1100710816)
    arr = -0.5 + rnd.random(N).astype(dtype)
    assert_arg_sorted(arr, np.argsort(arr, kind='quick'))

@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
def test_gh_22683():
    b = 777.68760986
    a = np.array([b] * 10000, dtype=object)
    refc_start = sys.getrefcount(b)
    np.choose(np.zeros(10000, dtype=int), [a], out=a)
    np.choose(np.zeros(10000, dtype=int), [a], out=a)
    refc_end = sys.getrefcount(b)
    assert refc_end - refc_start < 10


def test_gh_24459():
    a = np.zeros((50, 3), dtype=np.float64)
    with pytest.raises(TypeError):
        np.choose(a, [3, -1])


def test_gh_28206():
    a = np.arange(3)
    b = np.ones((3, 3), dtype=np.int64)
    out = np.array([np.nan, np.nan, np.nan])

    with warnings.catch_warnings():