Function test_interpolate_chunk_1d

xarray/tests/test_interp.py:903–966

Interpolate nd array with multiple independent indexers It should do a series of 1d interpolation

(
    method: InterpOptions, data_ndim, interp_ndim, nscalar, chunked: bool
)

)
@pytest.mark.filterwarnings("ignore:Increasing number of chunks")
def test_interpolate_chunk_1d(
    method: InterpOptions, data_ndim, interp_ndim, nscalar, chunked: bool
) -> None:
    """Interpolate nd array with multiple independent indexers

    It should do a series of 1d interpolation
    """

    if method in ["cubic", "pchip", "quintic"] and interp_ndim == 3:
        pytest.skip("Too slow.")

    # 3d non chunked data
    x = np.linspace(0, 1, 6)
    y = np.linspace(2, 4, 7)
    z = np.linspace(-0.5, 0.5, 8)
    da = xr.DataArray(
        data=np.sin(x[:, np.newaxis, np.newaxis])
        * np.cos(y[:, np.newaxis])
        * np.exp(z),
        coords=[("x", x), ("y", y), ("z", z)],
    )

    # choose the data dimensions
    for data_dims in permutations(da.dims, data_ndim):
        # select only data_ndim dim
        da = da.isel(  # take the middle line
            {dim: len(da.coords[dim]) // 2 for dim in da.dims if dim not in data_dims}
        )

        # chunk data
        da = da.chunk(chunks={dim: i + 1 for i, dim in enumerate(da.dims)})

        # choose the interpolation dimensions
        for interp_dims in permutations(da.dims, interp_ndim):
            # choose the scalar interpolation dimensions
            for scalar_dims in combinations(interp_dims, nscalar):
                dest = {}
                for dim in interp_dims:
                    if dim in scalar_dims:
                        # take the middle point
                        dest[dim] = 0.5 * (da.coords[dim][0] + da.coords[dim][-1])
                    else:
                        # pick some points, including outside the domain
                        before = 2 * da.coords[dim][0] - da.coords[dim][1]
                        after = 2 * da.coords[dim][-1] - da.coords[dim][-2]

                        dest[dim] = cast(
                            xr.DataArray,
                            np.linspace(
                                before.item(), after.item(), len(da.coords[dim]) * 13
                            ),
                        )
                        if chunked:
                            dest[dim] = xr.DataArray(data=dest[dim], dims=[dim])
                            dest[dim] = dest[dim].chunk(2)
                actual = da.interp(method=method, **dest)
                expected = da.compute().interp(method=method, **dest)