BUG: Squeeze when using interpolate_na with extra dim

[snowman2]

Possibly related to: pydata/xarray#9372

https://github.com/corteva/rioxarray/actions/runs/11151510847/job/30995110503

_______________ test_interpolate_na__all_nodata[open_rasterio] ________________
interpolate_na_nan = {'compare': '/home/runner/work/rioxarray/rioxarray/test/test_data/compare/MODIS_ARRAY_INTERPOLATE_NAN.nc', 'input': '/...nner/work/rioxarray/rioxarray/test/test_data/input/MODIS_ARRAY.nc', 'open': <function open_rasterio at 0x7f3fefa3b4c0>}
def test_interpolate_na__all_nodata(interpolate_na_nan):
rio_opened = "open_rasterio " in str(interpolate_na_nan["open"])
with interpolate_na_nan["open"](
interpolate_na_nan["input"], mask_and_scale=True
) as mda, interpolate_na_nan["open"](
interpolate_na_nan["compare"], mask_and_scale=True
) as mdc:
if hasattr(mda, "variables"):
for var in mda.rio.vars:
mda[var].values[~numpy.isnan(mda[var].values)] = numpy.nan
else:
mda.values[~numpy.isnan(mda.values)] = numpy.nan
> interpolated_ds = mda.rio.interpolate_na()
test/integration/test_integration_rioxarray.py:1517: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ rioxarray/raster_array.py:1063: in interpolate_na
interp_array = xarray.DataArray(
../../../micromamba/envs/test/lib/python3.11/site-packages/xarray/core/dataarray.py:479: in __init__
coords, dims = _infer_coords_and_dims(data.shape, coords, dims)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ shape = (1, 1, 200, 200)
coords = Coordinates:
* band (band) int64 8B 1
* x (x) float64 2kB -7.274e+06 -7.274e+06 ... -7.228e+06 -7.228e+06
* y (y) float64 2kB 5.05e+06 5.05e+06 ... 5.004e+06 5.004e+06
spatial_ref int64 8B 0
dims = ('band', 'y', 'x')
def _infer_coords_and_dims(
shape: tuple[int, ...],
coords: (
Sequence[Sequence | pd.Index | DataArray | Variable | np.ndarray]
| Mapping
| None
),
dims: str | Iterable[Hashable] | None,
) -> tuple[Mapping[Hashable, Any], tuple[Hashable, ...]]:
"""All the logic for creating a new DataArray"""
if (
coords is not None
and not utils.is_dict_like(coords)
and len(coords) != len(shape)
):
raise ValueError(
f"coords is not dict-like, but it has {len(coords)} items, "
f"which does not match the {len(shape)} dimensions of the "
"data"
)
if isinstance(dims, str):
dims = (dims,)
elif dims is None:
dims = [f"dim_{n}" for n in range(len(shape))]
if coords is not None and len(coords) == len(shape):
# try to infer dimensions from coords
if utils.is_dict_like(coords):
dims = list(coords.keys())
else:
for n, (dim, coord) in enumerate(zip(dims, coords, strict=True)):
coord = as_variable(
coord, name=dims[n], auto_convert=False
).to_index_variable()
dims[n] = coord.name
dims_tuple = tuple(dims)
if len(dims_tuple) != len(shape):
> raise ValueError(
"different number of dimensions on data "
f"and dims: {len(shape)} vs {len(dims_tuple)}"
)
E ValueError: different number of dimensions on data and dims: 4 vs 3
../../../micromamba/envs/test/lib/python3.11/site-packages/xarray/core/dataarray.py:185: ValueError
_____________ test_interpolate_na__all_nodata[interpolate_na_nan3] _____________
interpolate_na_nan = {'compare': '/home/runner/work/rioxarray/rioxarray/test/test_data/compare/MODIS_ARRAY_INTERPOLATE_NAN.nc', 'input': '/...st_data/input/MODIS_ARRAY.nc', 'open': functools.partial(<function open_dataset at 0x7f3fefd97880>, engine='rasterio')}
def test_interpolate_na__all_nodata(interpolate_na_nan):
rio_opened = "open_rasterio " in str(interpolate_na_nan["open"])
with interpolate_na_nan["open"](
interpolate_na_nan["input"], mask_and_scale=True
) as mda, interpolate_na_nan["open"](
interpolate_na_nan["compare"], mask_and_scale=True
) as mdc:
if hasattr(mda, "variables"):
for var in mda.rio.vars:
mda[var].values[~numpy.isnan(mda[var].values)] = numpy.nan
else:
mda.values[~numpy.isnan(mda.values)] = numpy.nan
> interpolated_ds = mda.rio.interpolate_na()
test/integration/test_integration_rioxarray.py:1517: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ rioxarray/raster_dataset.py:441: in interpolate_na
.rio.interpolate_na(method=method)
rioxarray/raster_array.py:1063: in interpolate_na
interp_array = xarray.DataArray(
../../../micromamba/envs/test/lib/python3.11/site-packages/xarray/core/dataarray.py:479: in __init__
coords, dims = _infer_coords_and_dims(data.shape, coords, dims)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ shape = (1, 1, 200, 200)
coords = Coordinates:
* band (band) int64 8B 1
* x (x) float64 2kB -7.274e+06 -7.274e+06 ... -7.228e+06 -7.228e+06
* y (y) float64 2kB 5.05e+06 5.05e+06 ... 5.004e+06 5.004e+06
spatial_ref int64 8B 0
dims = ('band', 'y', 'x')
def _infer_coords_and_dims(
shape: tuple[int, ...],
coords: (
Sequence[Sequence | pd.Index | DataArray | Variable | np.ndarray]
| Mapping
| None
),
dims: str | Iterable[Hashable] | None,
) -> tuple[Mapping[Hashable, Any], tuple[Hashable, ...]]:
"""All the logic for creating a new DataArray"""
if (
coords is not None
and not utils.is_dict_like(coords)
and len(coords) != len(shape)
):
raise ValueError(
f"coords is not dict-like, but it has {len(coords)} items, "
f"which does not match the {len(shape)} dimensions of the "
"data"
)
if isinstance(dims, str):
dims = (dims,)
elif dims is None:
dims = [f"dim_{n}" for n in range(len(shape))]
if coords is not None and len(coords) == len(shape):
# try to infer dimensions from coords
if utils.is_dict_like(coords):
dims = list(coords.keys())
else:
for n, (dim, coord) in enumerate(zip(dims, coords, strict=True)):
coord = as_variable(
coord, name=dims[n], auto_convert=False
).to_index_variable()
dims[n] = coord.name
dims_tuple = tuple(dims)
if len(dims_tuple) != len(shape):
> raise ValueError(
"different number of dimensions on data "
f"and dims: {len(shape)} vs {len(dims_tuple)}"
)
E ValueError: different number of dimensions on data and dims: 4 vs 3
../../../micromamba/envs/test/lib/python3.11/site-packages/xarray/core/dataarray.py:185: ValueError

[snowman2]

Thanks @justingruca 👍