How can I unstack a DataArray after having reassigned the MultiIndex?

[nicrie]

I just came across this unexpected behaviour (at least for me ;-) ) regarding unstacking a DataArray. What I want to do, in a nutshell, is to

1. stack a DataArray
2. replace the new MultiIndex with a 1D index
3. <... perform some operations on the DataArray that do not affect the MultiIndex dimension...>
4. Reassign the original MultiIndex.
5. unstack the DataArray to get the original form.

The last operation somehow fails (it does nothing). I would like to understand why and especially how to achieve this last step? The problem seems to be related to #5989, but as far as I can see, I do what is suggested there as a solution.

Minimal working example

import xarray as xr

da = xr.DataArray(
  [[1, 2.], [3, 4]],
  dims=['lat', 'lon'],
  coords=dict(lon=[5., 10], lat=[10., 12])
)

# 1. stack a `DataArray`
da_stacked = da.stack(x=('lat', 'lon'))
stacked_coords = da_stacked.coords['x']

# 2. replace the new `MultiIndex` with a 1D index
new_coords = np.arange(stacked_coords.size)
da_stacked = da_stacked.assign_coords(x=new_coords)

# 3. <...do some operations...>
# nothing happens

# 4. reassign the original `MultiIndex`
da_stacked = da_stacked.assign_coords(x=stacked_coords)

# 5. unstack the `DataArray` to obtain the original shape.
da_stacked.unstack()

Result of the last line of code

>>> xarray.DataArray (x: 4)
  array([1., 2., 3., 4.])
  Coordinates:
    x (x) object MultiIndex
    lat (x) float64 10.0 10.0 12.0 12.0
    lon (x) float64 5.0 10.0 5.0 10.0
  Attributes: (0)

What I actually expected: the input DataArray:

>>> xarray.DataArray (lat: 2, lon: 2)
  array([[1., 2.],
         [3., 4.]])
  Coordinates:
    lon (lon) float64 5.0 10.0
    lat (lat) float64 10.0 12.0
  Attributes: (0)

Additional info

xr.show_versions()

INSTALLED VERSIONS
------------------
commit: None
python: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:58:50) 
[GCC 10.3.0]
python-bits: 64
OS: Linux
OS-release: 5.15.0-53-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.2
libnetcdf: 4.8.1

xarray: 2022.6.0
pandas: 1.4.4
numpy: 1.23.2
scipy: 1.9.1
netCDF4: 1.6.0
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: None
cftime: 1.6.1
nc_time_axis: 1.4.1
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: 1.3.5
dask: 2022.9.0
distributed: 2022.9.0
matplotlib: 3.5.3
cartopy: 0.20.3
seaborn: 0.12.0
numbagg: None
fsspec: 2022.8.2
cupy: None
pint: None
sparse: None
flox: None
numpy_groupies: None
setuptools: 65.3.0
pip: 22.2.2
conda: None
pytest: None
IPython: 8.5.0
sphinx: 5.1.1

[mathause]

In step 4 you can do data = data.set_index(x=('lat', 'lon')) and then data.unstack("x")

[nicrie]

Thank you, that worked indeed! But I'm curious - why does it work like this? What's happening with the DataArray when I set the index like this? There must be a more fundamental difference between coords and index that I don't fully understand...

Interestingly, trying my example above in a slightly older version of xarray (version: 0.21.1) it actually worked without explicitly setting the index. Just for me to understand, can you tell me what changed in the way how coordinates and indices work?

Thanks in advance!

[TomNicholas]

There have been some very big changes to how xarray handles indexes internally recently. (See our roadmap and the corresponding project.) I'm not sure exactly what changed in this case but @benbovy might be able to enlighten us. 🙂

[benbovy]

This is because in the future replacing some level and/or dimension coordinate(s) of a multi-index with assign_coords() will drop the index. As unstack requires a multi-index, you first need to re-build it explicitly from the new coordinates using set_index.

The reason is that coordinates and indexes are now more decoupled concepts in the Xarray data model: one index may correspond to multiple coordinates: multi-indexes no more have "virtual" level coordinates. When replacing only some of the coordinates having the same index, it is easier to just drop the index instead of trying to handle all the special cases that may arise in this situation.

in version 2022.6.0 (the one your are using) the change has actually been introduced without deprecation notice, and with some weird behavior like in your example. This has been reverted in #6798 and #7101, so with the latest version this should work again and you should see a deprecation warning.

[nicrie]

Thanks to both of you, Tom and Benoit! Really appreciate your (ultra fast) feedback and the provided references.