extrapolate not working for multi-dimentional data

[zxdawn]

What happened?

I'm trying to interpolate multi-dimensional data. But, it doesn't work with extrapolating.

What did you expect to happen?

No response

Minimal Complete Verifiable Example

import xarray as xr
import numpy as np

da = xr.DataArray(
        np.sin(0.3 * np.arange(20).reshape(5, 4)),
        [("x", np.arange(5)), ("y", [0.1, 0.2, 0.3, 0.4])],
)
x = xr.DataArray([-0.5, 1.5, 2.5], dims="z")

y = xr.DataArray([0.15, 0.25, 0.35], dims="z")

da.interp(x=x, y=y, kwargs={"fill_value": "extrapolate"})

Relevant log output

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [2], in <module>
      8 x = xr.DataArray([-0.5, 1.5, 2.5], dims="z")
     10 y = xr.DataArray([0.15, 0.25, 0.35], dims="z")
---> 12 da.interp(x=x, y=y, kwargs={"fill_value": "extrapolate"})

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/dataarray.py:1742, in DataArray.interp(self, coords, method, assume_sorted, kwargs, **coords_kwargs)
   1737 if self.dtype.kind not in "uifc":
   1738     raise TypeError(
   1739         "interp only works for a numeric type array. "
   1740         "Given {}.".format(self.dtype)
   1741     )
-> 1742 ds = self._to_temp_dataset().interp(
   1743     coords,
   1744     method=method,
   1745     kwargs=kwargs,
   1746     assume_sorted=assume_sorted,
   1747     **coords_kwargs,
   1748 )
   1749 return self._from_temp_dataset(ds)

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/dataset.py:3192, in Dataset.interp(self, coords, method, assume_sorted, kwargs, method_non_numeric, **coords_kwargs)
   3189 if dtype_kind in "uifc":
   3190     # For normal number types do the interpolation:
   3191     var_indexers = {k: v for k, v in use_indexers.items() if k in var.dims}
-> 3192     variables[name] = missing.interp(var, var_indexers, method, **kwargs)
   3193 elif dtype_kind in "ObU" and (use_indexers.keys() & var.dims):
   3194     # For types that we do not understand do stepwise
   3195     # interpolation to avoid modifying the elements.
   (...)
   3198     # this loop there might be some duplicate code that slows it
   3199     # down, therefore collect these signals and run it later:
   3200     to_reindex[name] = var

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/missing.py:640, in interp(var, indexes_coords, method, **kwargs)
    638 original_dims = broadcast_dims + dims
    639 new_dims = broadcast_dims + list(destination[0].dims)
--> 640 interped = interp_func(
    641     var.transpose(*original_dims).data, x, destination, method, kwargs
    642 )
    644 result = Variable(new_dims, interped, attrs=var.attrs)
    646 # dimension of the output array

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/missing.py:765, in interp_func(var, x, new_x, method, kwargs)
    749         meta = var._meta
    751     return da.blockwise(
    752         _dask_aware_interpnd,
    753         out_ind,
   (...)
    762         align_arrays=False,
    763     )
--> 765 return _interpnd(var, x, new_x, func, kwargs)

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/missing.py:789, in _interpnd(var, x, new_x, func, kwargs)
    787 # stack new_x to 1 vector, with reshape
    788 xi = np.stack([x1.values.ravel() for x1 in new_x], axis=-1)
--> 789 rslt = func(x, var, xi, **kwargs)
    790 # move back the interpolation axes to the last position
    791 rslt = rslt.transpose(range(-rslt.ndim + 1, 1))

File ~/miniconda3/lib/python3.9/site-packages/scipy/interpolate/interpolate.py:2703, in interpn(points, values, xi, method, bounds_error, fill_value)
   2701 # perform interpolation
   2702 if method == "linear":
-> 2703     interp = RegularGridInterpolator(points, values, method="linear",
   2704                                      bounds_error=bounds_error,
   2705                                      fill_value=fill_value)
   2706     return interp(xi)
   2707 elif method == "nearest":

File ~/miniconda3/lib/python3.9/site-packages/scipy/interpolate/interpolate.py:2469, in RegularGridInterpolator.__init__(self, points, values, method, bounds_error, fill_value)
   2465     fill_value_dtype = np.asarray(fill_value).dtype
   2466     if (hasattr(values, 'dtype') and not
   2467             np.can_cast(fill_value_dtype, values.dtype,
   2468                         casting='same_kind')):
-> 2469         raise ValueError("fill_value must be either 'None' or "
   2470                          "of a type compatible with values")
   2472 for i, p in enumerate(points):
   2473     if not np.all(np.diff(p) > 0.):

ValueError: fill_value must be either 'None' or of a type compatible with values

Anything else we need to know?

I've also tried da.interp(x=x, kwargs={"fill_value": "extrapolate"}).interp(y=y), but I got this error:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Input In [4], in <module>
     10 y = xr.DataArray([0.15, 0.25, 0.35], dims="z")
     12 # da.interp(x=x, y=y, kwargs={"fill_value": "extrapolate"})
---> 13 da.interp(x=x, kwargs={"fill_value": "extrapolate"}).interp(y=y)

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/dataarray.py:1742, in DataArray.interp(self, coords, method, assume_sorted, kwargs, **coords_kwargs)
   1737 if self.dtype.kind not in "uifc":
   1738     raise TypeError(
   1739         "interp only works for a numeric type array. "
   1740         "Given {}.".format(self.dtype)
   1741     )
-> 1742 ds = self._to_temp_dataset().interp(
   1743     coords,
   1744     method=method,
   1745     kwargs=kwargs,
   1746     assume_sorted=assume_sorted,
   1747     **coords_kwargs,
   1748 )
   1749 return self._from_temp_dataset(ds)

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/dataset.py:3130, in Dataset.interp(self, coords, method, assume_sorted, kwargs, method_non_numeric, **coords_kwargs)
   3121 if coords:
   3122     # This avoids broadcasting over coordinates that are both in
   3123     # the original array AND in the indexing array. It essentially
   3124     # forces interpolation along the shared coordinates.
   3125     sdims = (
   3126         set(self.dims)
   3127         .intersection(*[set(nx.dims) for nx in indexers.values()])
   3128         .difference(coords.keys())
   3129     )
-> 3130     indexers.update({d: self.variables[d] for d in sdims})
   3132 obj = self if assume_sorted else self.sortby([k for k in coords])
   3134 def maybe_variable(obj, k):
   3135     # workaround to get variable for dimension without coordinate.

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/dataset.py:3130, in <dictcomp>(.0)
   3121 if coords:
   3122     # This avoids broadcasting over coordinates that are both in
   3123     # the original array AND in the indexing array. It essentially
   3124     # forces interpolation along the shared coordinates.
   3125     sdims = (
   3126         set(self.dims)
   3127         .intersection(*[set(nx.dims) for nx in indexers.values()])
   3128         .difference(coords.keys())
   3129     )
-> 3130     indexers.update({d: self.variables[d] for d in sdims})
   3132 obj = self if assume_sorted else self.sortby([k for k in coords])
   3134 def maybe_variable(obj, k):
   3135     # workaround to get variable for dimension without coordinate.

File ~/miniconda3/lib/python3.9/site-packages/xarray/core/utils.py:459, in Frozen.__getitem__(self, key)
    458 def __getitem__(self, key: K) -> V:
--> 459     return self.mapping[key]

KeyError: 'z'

Environment

INSTALLED VERSIONS
------------------
commit: None
python: 3.9.7 | packaged by conda-forge | (default, Sep 29 2021, 19:20:46) 
[GCC 9.4.0]
python-bits: 64
OS: Linux
OS-release: 5.11.0-40-generic
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.12.1
libnetcdf: 4.8.1

xarray: 0.20.2
pandas: 1.3.5
numpy: 1.19.5
scipy: 1.7.3
netCDF4: 1.5.8
pydap: None
h5netcdf: None
h5py: 3.6.0
Nio: None
zarr: 2.10.3
cftime: 1.5.1.1
nc_time_axis: None
PseudoNetCDF: None
rasterio: 1.2.10
cfgrib: None
iris: None
bottleneck: None
dask: 2021.12.0
distributed: 2021.12.0
matplotlib: 3.5.1
cartopy: 0.20.2
seaborn: None
numbagg: None
fsspec: 2022.01.0
cupy: None
pint: 0.18
sparse: None
setuptools: 59.8.0
pip: 21.3.1
conda: 4.11.0
pytest: None
IPython: 8.0.0
sphinx: None

[andersy005]

@zxdawn,

Thank you for the thorough, reproducible example.

Per Scipy's documentation, RegularGridInterpolator uses fill_value=None for extrapolation

In [23]: da.interp(x=x, y=y, kwargs={"fill_value": None})

Out[23]: 
<xarray.DataArray (z: 3)>
array([-0.26074336,  0.63496063, -0.46643289])
Coordinates:
    x        (z) float64 -0.5 1.5 2.5
    y        (z) float64 0.15 0.25 0.35
Dimensions without coordinates: z

[zxdawn]

@andersy005 I tested with another example:

import xarray as xr
import numpy as np

da = xr.DataArray(np.random.rand(5, 5), dims=['x','y'], coords=[range(5), range(5)])
da.interp(x=5, kwargs={"fill_value": None})

The result is all nan values without extrapolation. Do you know how to deal with it?

The 1d extrapolation works well:

da = xr.DataArray(np.random.rand(5), dims=['x'], coords=[range(5)])
da.interp(x=5, kwargs={"fill_value": 'extrapolate'})

>>> array(0.46080927)

[zxdawn]

Strange ... This multi-dimension interpolation seems to work well:

import xarray as xr
import numpy as np

da = xr.DataArray(np.random.rand(5, 5, 5, 5), dims=['x','y', 'var1', 'var2'], coords=[range(5)]*4)
var1 = xr.DataArray(np.random.rand(5, 5)-10, dims=['x','y'])
var2 = xr.DataArray(np.random.rand(5, 5)-10, dims=['x','y'])

da.interp(var1=var1, var2=var2, kwargs={"fill_value": None})

[zxdawn]

@andersy005 Thanks a lot, I realize that it's mentioned in the comments of Guide.

[rafa-guedes]

Extrapolation does not seem to work when trying to interpolate onto the existing axes rather than onto the new "z" axis as in the first example:

xi = np.arange(-0.5, 5.0, 0.5)
yi = np.arange(0, 0.6, 0.1)
dai = da.interp(x=xi, y=yi, kwargs={"fill_value": None})

In [48]: scipy.__version__
Out[48]: '1.8.0'

In [49]: np.__version__
Out[49]: '1.22.3'

In [50]: xr.__version__
Out[50]: '2022.3.0'

See figures below for the original and interpolated data. I was sure this worked before so not sure something changed in the interpolation libraries, or if the kwargs are not propagating through in this case (the extrapolation does work when running this directly with scipy.interpolate.interpn and passing fill_value=None along with bounds_error=False).