unstack() appears to fail to unstack a MultiIndex under specific circumstances. What's going on?

[DaniJonesOcean]

Hi all. First, let me say that it's entirely possible that I'm simply doing something wrong. Any advice/help is appreciated!

Problem description

I'm working with some gridded data (i.e. with two dimensions, x and y), but at times I need to treat the gridded data with a single index (i.e. profile number or z). This problem arises when I try to convert between the two approaches.

Here's what I'm trying to do, specifically: I need to combine a few 1D DataArrays into an existing Dataset that uses a MultiIndex. The idea is to combine all the 1D DataArrays together and then use unstack() to recover the original latitude and longitude coordinates. This seems to work okay when I combine my existing Dataset with a single DataArray, but it fails when I assign multiple new DataArrays to my existing Dataset.

What happened:

The unstack() command doesn't seem to do anything. It returns an unchanged Dataset with a single dimension.

What you expected to happen:

I expected the unstack() command to return a Dataset with two dimensions. Interestingly, this is exactly what it does if I only assign a single new DataArray to my existing Dataset. The problem occurs when I add more than one DataArray. The code below reproduces the problem.

Try it out: if you comment out one of the lines indicated below, then the code returns a two-dimensional Dataset.

Minimal Complete Verifiable Example:

# import packages
import xarray as xr
import numpy as np

# generate random example data on a 2D "grid"
np.random.seed(0)
temperature = 15 + 8 * np.random.randn(2, 2)
precipitation = 10 * np.random.rand(2, 2)
lon = [[-99.83, -99.32], [-99.79, -99.23]]
lat = [[42.25, 42.21], [42.63, 42.59]]

# create DataArray using the "temperature" data
ds = xr.Dataset(
    data_vars=dict(
        temperature=(["x", "y"], temperature),
        precipitation=(["x", "y"], precipitation),
    ),
    coords=dict(
        lon=(["x", "y"], lon),
        lat=(["x", "y"], lat),
    ),
    attrs=dict(description="Weather related data."),
)

# stack the Dataset
ds_stacked = ds.stack(z=("x","y"))
ds_stacked

# create 1D array of labels (to be associated with each lat/lon pair)
d1 = xr.DataArray(
    data=np.random.randint(low=0, high=2, size=(4,)),
    dims=["z"],
    attrs=dict(
        description="Label",
    ),
)

# create 1D array of labels (to be associated with each lat/lon pair)
d2 = xr.DataArray(
    data=np.random.randn(4,),
    dims=["z"],
    attrs=dict(
        description="Posteriors",
    ),
)

# assign the labels in d1 to da_stacked
ds_stacked['label'] = d1 #<-- comment out either this line...
ds_stacked['post'] = d2 #<-- or this line...and the unstack() command does what I expect it to do! 
ds_stacked

# finally, unstack 
ds_final = ds_stacked.unstack() # <-- this should unstack the "z" MultiIndex coordinate.
ds_final 

Anything else we need to know?:

I'm reasonably convinced that this method will do what I need it to do, if I can get it to work. That being said, I'm open to other alternative approaches to this problem. : )

Any thoughts or suggestions?

Environment:

Output of xr.show_versions()

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.10.47-linuxkit
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: C.UTF-8
LANG: C.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.12.1
libnetcdf: 4.8.1

xarray: 0.20.1
pandas: 1.3.4
numpy: 1.21.4
scipy: 1.7.2
netCDF4: 1.5.8
pydap: installed
h5netcdf: 0.11.0
h5py: 3.4.0
Nio: None
zarr: 2.10.1
cftime: 1.5.1.1
nc_time_axis: 1.4.0
PseudoNetCDF: None
rasterio: 1.2.10
cfgrib: 0.9.9.1
iris: None
bottleneck: 1.3.2
dask: 2021.11.1
distributed: 2021.11.1
matplotlib: 3.4.3
cartopy: 0.20.1
seaborn: 0.11.2
numbagg: None
fsspec: 2021.11.0
cupy: None
pint: 0.18
sparse: 0.13.0
setuptools: 58.5.3
pip: 20.3.4
conda: None
pytest: None
IPython: 7.29.0
sphinx: None

[mathause]

Assigning d1 or d1 your ds_stacked looses the MultiIndex compare

Before:

<xarray.Dataset>
Dimensions:        (z: 4)
Coordinates:
  * z              (z) MultiIndex
  - x              (z) int64 0 0 1 1
  - y              (z) int64 0 1 0 1

After:

<xarray.Dataset>
Dimensions:        (z: 4)
Coordinates:
  * z              (z) object (0, 0) (0, 1) (1, 0) (1, 1)
    lon            (z) float64 -99.83 -99.32 -99.79 -99.23
    lat            (z) float64 42.25 42.21 42.63 42.59

I am not entirely why this happens... But to work around this issue you have two possibilities:

# assign a raw array
ds_stacked['label'] = ("z", d1.data)

# add the correct coords
d2 = d2.assign_coords(z=ds_stacked.z)
ds_stacked['post'] = d2