Function _interpnd

xarray/core/missing.py:786–843

Core nD array interpolation routine. The first half arrays in `coords` are original coordinates, the other half are destination coordinates.

(
    data: np.ndarray,
    *coords: np.ndarray,
    interp_func: Interpolator | InterpCallable,
    interp_kwargs,
    result_coord_core_dims: list[tuple[Hashable, ...]],
)

def _interpnd(
    data: np.ndarray,
    *coords: np.ndarray,
    interp_func: Interpolator | InterpCallable,
    interp_kwargs,
    result_coord_core_dims: list[tuple[Hashable, ...]],
) -> np.ndarray:
    """
    Core nD array interpolation routine.
    The first half arrays in `coords` are original coordinates,
    the other half are destination coordinates.
    """
    n_x = len(coords) // 2
    ndim = data.ndim
    nconst = ndim - n_x

    # Convert everything to Variables, since that makes applying
    # `_localize` and `_floatize_x` much easier
    x = [
        Variable([f"dim_{nconst + dim}"], _x, fastpath=True)
        for dim, _x in enumerate(coords[:n_x])
    ]
    new_x = list(
        broadcast_variables(
            *(
                Variable(dims, _x, fastpath=True)
                for dims, _x in zip(result_coord_core_dims, coords[n_x:], strict=True)
            )
        )
    )
    var = Variable([f"dim_{dim}" for dim in range(ndim)], data, fastpath=True)

    if interp_kwargs.get("method") in ["linear", "nearest"]:
        indexes_coords = {
            _x.dims[0]: (_x, _new_x) for _x, _new_x in zip(x, new_x, strict=True)
        }
        # simple speed up for the local interpolation
        var, indexes_coords = _localize(var, indexes_coords)
        x, new_x = tuple(
            list(_)
            for _ in zip(*(indexes_coords[d] for d in indexes_coords), strict=True)
        )

    x_list, new_x_list = _floatize_x(x, new_x)

    if len(x) == 1:
        # TODO: narrow interp_func to interpolator here
        return _interp1d(var, x_list, new_x_list, interp_func, interp_kwargs)  # type: ignore[arg-type]

    # move the interpolation axes to the start position
    data = transpose(var._data, range(-len(x), var.ndim - len(x)))

    # stack new_x to 1 vector, with reshape
    xi = stack([ravel(x1.data) for x1 in new_x_list], axis=-1)
    rslt: np.ndarray = interp_func(x_list, data, xi, **interp_kwargs)  # type: ignore[assignment]
    # move back the interpolation axes to the last position
    rslt = transpose(rslt, range(-rslt.ndim + 1, 1))
    return reshape(rslt, rslt.shape[:-1] + new_x[0].shape)