Stacking and Unstacking DataArray with 1D MultiIndex

[nicrie]

I'm working on a project that includes a process where I need to:

1. Stack a DataArray
2. Substitute the new MultiIndex with a 1D index
3. Execute some operations on the DataArray that don't affect the MultiIndex dimension
4. Restore the original MultiIndex
5. Unstack the DataArray to regain the original dimension/coordinates

Now, here's the caveat: If the MultiIndex is already a 1D index, the solution I've employed earlier (#7337 ) doesn't hold water.

For instance, when I generate a 3D DataArray with dimensions (time, lat, lon), stack it into a 2D version with dimensions (sample, feature) where sample = ('time', ) and feature = ('lat', 'lon'), a problem arises. After carrying out the operations, I reassign the original MultiIndex and unstack. While unstacking works well for feature, it fails for sample. The final output dimensions are (sample, lon, lat), but I need them to be (time, lon, lat).

I understand this might seem a bit unusual, but it is a part of a more extensive procedure. I know I could just write some code to handle this 1D MultiIndex case and rename instead of stack, and I will if there's no other way. However, I'm curious if anyone could shed light on the behaviour I'm currently observing.

Below is a minimal working example demonstrating the issue:

MWE

# Create a 3D DataArray
data = xr.DataArray(
    np.random.rand(5, 2, 2),
    dims=['time', 'lon', 'lat'],
    coords={'time': np.arange(5), 'lon': np.arange(2), 'lat': np.arange(2)}
)

# Stack DataArray into a 2D version with dimension (sample x feature)
# Samples = time, Features = (lon, lat)
data_stacked = data.stack(sample=('time',), feature=('lon', 'lat'))
original_coords = data_stacked.coords
data_stacked = data_stacked.drop_vars({'sample'}).assign_coords(sample=np.arange(5))
data_stacked = data_stacked.drop_vars({'feature'}).assign_coords(feature=np.arange(4))

# ... do some operations on data_stacked

data_stacked = data_stacked.assign_coords(feature=original_coords['feature']).set_index(feature=['lon', 'lat'])
data_stacked = data_stacked.assign_coords(sample=original_coords['sample']).set_index(sample=('time',))

data_reconstruced = data_stacked.unstack()

assert data.dims == data_reconstruced.dims  # raises an AssertionError

The assertion at the end raises an AssertionError. How can I fix this issue and ensure the dimensions after unstacking are the same as the original data?