def test_lazy_indexing(use_dask):
idx = xr.DataArray([[0, 10], [33, 99]], dims=("x", "y"))
da = xr.DataArray(np.arange(100), dims=("time", ))
db = xr.DataArray(-np.arange(100), dims=("time", ))
if use_dask:
idx = idx.chunk({"x": 1})
# Ensure tasks are different
outa, outb = compute(rl.lazy_indexing(da, idx), rl.lazy_indexing(db, idx))
assert set(outa.dims) == {"x", "y"}
np.testing.assert_array_equal(idx, outa)
np.testing.assert_array_equal(idx, -outb)
assert "time" not in outa.coords
def test_lazy_indexing_special_cases(use_dask):
a = xr.DataArray(np.random.rand(10, 10, 10), dims=("x", "y", "z"))
b = xr.DataArray(np.random.rand(10, 10, 10), dims=("x", "y", "z"))
if use_dask:
a = a.chunk({"y": 5, "z": 5})
b = b.chunk({"y": 5, "z": 1})
with pytest.raises(ValueError):
rl.lazy_indexing(a, b)
b = b.argmin("x").argmin("y")
with pytest.raises(ValueError,
match="more than one dimension more than index"):
rl.lazy_indexing(a, b)