Repository: dask/dask
File: dask/array/tests/test_slicing.py (Python)

Code before:
    index2 = np.array([0, 2, 4, 6, 1, 3, 5, 7])
    index3 = np.array([3, 0, 2, 1, 7, 4, 6, 5])
    np.testing.assert_array_equal(a, index2)
    np.testing.assert_array_equal(b, index3)


@pytest.mark.filterwarnings("ignore:Slicing:dask.array.core.PerformanceWarning")
@pytest.mark.parametrize(
    "size, chunks", [((100, 2), (50, 2)), ((100, 2), (37, 1)), ((100,), (55,))]
)
def test_shuffle_slice(size, chunks):
    x = da.random.default_rng().integers(0, 1000, size=size, chunks=chunks)
    index = np.arange(len(x))
    np.random.default_rng().shuffle(index)

    a = x[index]
    b = shuffle_slice(x, index)
    assert_eq(a, b)


def test_unknown_chunks_length_one():
    a = np.arange(256, dtype=int)
    arr = da.from_array(a, chunks=(256,))
    result = np.flatnonzero(arr)
    assert_eq(result[[0, -1]], np.flatnonzero(a)[[0, -1]])

    a = a.reshape(16, 16)
    arr = da.from_array(a, chunks=(8, 16))
    arr._chunks = ((8, 8), (np.nan,))
    result = arr[:, [0, -1]]
    expected = a[:, [0, -1]]
    assert_eq(result, expected)

    arr = da.from_array(a, chunks=(8, 8))
    arr._chunks = ((8, 8), (np.nan, np.nan))
    with pytest.raises(ValueError, match="Array chunk size or shape"):
        arr[:, [0, -1]]


@pytest.mark.parametrize("lock", [True, False])
@pytest.mark.parametrize("asarray", [True, False])
@pytest.mark.parametrize("fancy", [True, False])
def test_gh4043(lock, asarray, fancy):
    a1 = da.from_array(np.zeros(3), chunks=1, asarray=asarray, lock=lock, fancy=fancy)
    a2 = da.from_array(np.ones(3), chunks=1, asarray=asarray, lock=lock, fancy=fancy)
    al = da.stack([a1, a2])
    assert_eq(al, al)


def test_slice_array_3d_with_bool_numpy_array():
    # https://github.com/dask/dask/issues/6089

Review comment (suggestion):
```suggestion
    # np.flatnonzero dispatches
    result = np.flatnonzero(arr)
    assert_eq(result[[0, -1]], np.flatnonzero(a)[[0, -1]])
    
    result = da.flatnonzero(arr)
    assert_eq(result[[0, -1]], np.flatnonzero(a)[[0, -1]])
```

Code after:
    index2 = np.array([0, 2, 4, 6, 1, 3, 5, 7])
    index3 = np.array([3, 0, 2, 1, 7, 4, 6, 5])
    np.testing.assert_array_equal(a, index2)
    np.testing.assert_array_equal(b, index3)


@pytest.mark.filterwarnings("ignore:Slicing:dask.array.core.PerformanceWarning")
@pytest.mark.parametrize(
    "size, chunks", [((100, 2), (50, 2)), ((100, 2), (37, 1)), ((100,), (55,))]
)
def test_shuffle_slice(size, chunks):
    x = da.random.default_rng().integers(0, 1000, size=size, chunks=chunks)
    index = np.arange(len(x))
    np.random.default_rng().shuffle(index)

    a = x[index]
    b = shuffle_slice(x, index)
    assert_eq(a, b)


def test_unknown_chunks_length_one():
    a = np.arange(256, dtype=int)
    arr = da.from_array(a, chunks=(256,))
    # np.flatnonzero dispatches
    result = np.flatnonzero(arr)
    assert_eq(result[[0, -1]], np.flatnonzero(a)[[0, -1]])

    result = da.flatnonzero(arr)
    assert_eq(result[[0, -1]], np.flatnonzero(a)[[0, -1]])

    a = a.reshape(16, 16)
    arr = da.from_array(a, chunks=(8, 16))
    arr._chunks = ((8, 8), (np.nan,))
    result = arr[:, [0, -1]]
    expected = a[:, [0, -1]]
    assert_eq(result, expected)

    arr = da.from_array(a, chunks=(8, 8))
    arr._chunks = ((8, 8), (np.nan, np.nan))
    with pytest.raises(ValueError, match="Array chunk size or shape"):
        arr[:, [0, -1]]


@pytest.mark.parametrize("lock", [True, False])
@pytest.mark.parametrize("asarray", [True, False])
@pytest.mark.parametrize("fancy", [True, False])
def test_gh4043(lock, asarray, fancy):
    a1 = da.from_array(np.zeros(3), chunks=1, asarray=asarray, lock=lock, fancy=fancy)
    a2 = da.from_array(np.ones(3), chunks=1, asarray=asarray, lock=lock, fancy=fancy)
    al = da.stack([a1, a2])
    assert_eq(al, al)