Function test_map_blocks_dask_args

xarray/tests/test_dask.py:1258–1316

()

def test_map_blocks_dask_args():
    da1 = xr.DataArray(
        np.ones((10, 20)),
        dims=["x", "y"],
        coords={"x": np.arange(10), "y": np.arange(20)},
    ).chunk({"x": 5, "y": 4})

    # check that block shapes are the same
    def sumda(da1, da2):
        assert da1.shape == da2.shape
        return da1 + da2

    da2 = da1 + 1
    with raise_if_dask_computes():
        mapped = xr.map_blocks(sumda, da1, args=[da2])
    xr.testing.assert_equal(da1 + da2, mapped)

    # one dimension in common
    da2 = (da1 + 1).isel(x=1, drop=True)
    with raise_if_dask_computes():
        mapped = xr.map_blocks(operator.add, da1, args=[da2])
    xr.testing.assert_equal(da1 + da2, mapped)

    # test that everything works when dimension names are different
    da2 = (da1 + 1).isel(x=1, drop=True).rename({"y": "k"})
    with raise_if_dask_computes():
        mapped = xr.map_blocks(operator.add, da1, args=[da2])
    xr.testing.assert_equal(da1 + da2, mapped)

    with pytest.raises(ValueError, match=r"Chunk sizes along dimension 'x'"):
        xr.map_blocks(operator.add, da1, args=[da1.chunk({"x": 1})])

    with pytest.raises(ValueError, match=r"cannot align.*index.*are not equal"):
        xr.map_blocks(operator.add, da1, args=[da1.reindex(x=np.arange(20))])

    # reduction
    da1 = da1.chunk({"x": -1})
    da2 = da1 + 1
    with raise_if_dask_computes():
        mapped = xr.map_blocks(lambda a, b: (a + b).sum("x"), da1, args=[da2])
    xr.testing.assert_equal((da1 + da2).sum("x"), mapped)

    # reduction with template
    da1 = da1.chunk({"x": -1})
    da2 = da1 + 1
    with raise_if_dask_computes():
        mapped = xr.map_blocks(
            lambda a, b: (a + b).sum("x"), da1, args=[da2], template=da1.sum("x")
        )
    xr.testing.assert_equal((da1 + da2).sum("x"), mapped)

    # bad template: not chunked
    with pytest.raises(ValueError, match="Provided template has no dask arrays"):
        xr.map_blocks(
            lambda a, b: (a + b).sum("x"),
            da1,
            args=[da2],
            template=da1.sum("x").compute(),