Function test_interpolate_vectorize

xarray/tests/test_interp.py:204–290

(use_dask: bool, method: InterpOptions)

    ),
)
def test_interpolate_vectorize(use_dask: bool, method: InterpOptions) -> None:
    # scipy interpolation for the reference
    def func(obj, dim, new_x, method):
        scipy_kwargs = {}
        interpolant_options = {
            "barycentric": scipy.interpolate.BarycentricInterpolator,
            "krogh": scipy.interpolate.KroghInterpolator,
            "pchip": scipy.interpolate.PchipInterpolator,
            "akima": scipy.interpolate.Akima1DInterpolator,
            "makima": scipy.interpolate.Akima1DInterpolator,
        }

        shape = [s for i, s in enumerate(obj.shape) if i != obj.get_axis_num(dim)]
        for s in new_x.shape[::-1]:
            shape.insert(obj.get_axis_num(dim), s)

        if method in interpolant_options:
            interpolant = interpolant_options[method]
            if method == "makima":
                scipy_kwargs["method"] = method
            return interpolant(
                da[dim], obj.data, axis=obj.get_axis_num(dim), **scipy_kwargs
            )(new_x).reshape(shape)
        else:
            return scipy.interpolate.interp1d(
                da[dim],
                obj.data,
                axis=obj.get_axis_num(dim),
                kind=method,  # type: ignore[arg-type,unused-ignore]
                bounds_error=False,
                fill_value=np.nan,
                **scipy_kwargs,
            )(new_x).reshape(shape)

    da = get_example_data(0)

    if use_dask:
        da = da.chunk({"y": 5})

    # xdest is 1d but has different dimension
    xdest = xr.DataArray(
        np.linspace(0.1, 0.9, 30),
        dims="z",
        coords={"z": np.random.randn(30), "z2": ("z", np.random.randn(30))},
    )

    actual = da.interp(x=xdest, method=method)

    expected = xr.DataArray(
        func(da, "x", xdest, method),
        dims=["z", "y"],
        coords={
            "z": xdest["z"],
            "z2": xdest["z2"],
            "y": da["y"],
            "x": ("z", xdest.values),
            "x2": ("z", func(da["x2"], "x", xdest, method)),
        },