Unstacking the diagonals of a sequence of matrices raises ValueError: IndexVariable objects must be 1-dimensional

[DWesl]

What happened?

In my work, I produced a sequence of covariance matrices for a 2-D quantity. I wanted to extract the diagonal of the covariance matrices, then make that diagonal 2-D so I could plot it.

I could unstack the 2-D dimensions in the sequence of covariance matrices without issue. I figured out how to extract the diagonal of the covariance matrices. Unstacking the diagonal using the same procedure raised a ValueError.

What did you expect to happen?

I expected the sequence of one-dimensional diagonals to unstack into a sequence of two-dimensional fields so I could plot them with pcolormesh.

I can make this happen by unstacking the two dimensions (producing a 5-D DataArray) and extracting the diagonals from that, but I don't see a reason it shouldn't work in the other order.

Minimal Complete Verifiable Example

import numpy as np
import xarray

# Working:
test = xarray.DataArray(
    np.eye(12),
    dims=("dim0", "adj_dim0"),
    coords={
        "dim0_0": (("dim0",), np.repeat(np.arange(3), 4)),
        "dim0_1": (("dim0",), np.tile(np.arange(4), 3)),
    },
)
unstacked = test.set_index(dim0=["dim0_0", "dim0_1"]).unstack("dim0")
diag_index = xarray.DataArray(np.arange(test.shape[0]), dims=("diag",))
unstacked_diag = (
    test.isel(dim0=diag_index, adj_dim0=diag_index)
    .set_index(diag=["dim0_0", "dim0_1"])
    .unstack("diag")
)

# Not working:
test = xarray.DataArray(
    dims=("dim1", "dim0", "adj_dim0"),
    data=np.tile(np.eye(12), (2, 1, 1)),
    coords={
        "dim0": np.arange(12),
        "dim0_0": (("dim0",), np.repeat(np.arange(3), 4)),
        "dim0_1": (("dim0",), np.tile(np.arange(4), 3)),
        "adj_dim0": np.arange(12),
        "adj_dim0_0": (("adj_dim0",), np.repeat(np.arange(3), 4)),
        "adj_dim0_1": (("adj_dim0",), np.tile(np.arange(4), 3)),
    },
)
unstacked = test.set_index(
    dim0=["dim0_0", "dim0_1"], adj_dim0=["adj_dim0_0", "adj_dim0_1"]
).unstack(["dim0", "adj_dim0"])

diag_index0 = xarray.DataArray(np.arange(unstacked.shape[1]), dims=("diag_0",))
diag_index1 = xarray.DataArray(np.arange(unstacked.shape[2]), dims=("diag_1",))
unstacked_diag = unstacked.isel(
    dim0_0=diag_index0,
    dim0_1=diag_index1,
    adj_dim0_0=diag_index0,
    adj_dim0_1=diag_index1,
)

diag_index = xarray.DataArray(np.arange(test.shape[1]), dims=("diag",))
test.isel(dim0=diag_index, adj_dim0=diag_index).set_index(
    diag=["dim0_0", "dim0_1"]
).unstack("diag")

Relevant log output

Traceback (most recent call last):
  File "./test.py", line 50, in <module>
    ).unstack("diag")
  File "~/.conda/envs/plotting/lib/python3.10/site-packages/xarray/core/dataarray.py", line 2201, in unstack
    ds = self._to_temp_dataset().unstack(dim, fill_value, sparse)
  File "~/.conda/envs/plotting/lib/python3.10/site-packages/xarray/core/dataset.py", line 4214, in unstack
    result = result._unstack_once(dim, fill_value, sparse)
  File "~/.conda/envs/plotting/lib/python3.10/site-packages/xarray/core/dataset.py", line 4070, in _unstack_once
    variables[name] = var._unstack_once(
  File "~/.conda/envs/plotting/lib/python3.10/site-packages/xarray/core/variable.py", line 1690, in _unstack_once
    return self._replace(dims=new_dims, data=data)
  File "~/.conda/envs/plotting/lib/python3.10/site-packages/xarray/core/variable.py", line 978, in _replace
    return type(self)(dims, data, attrs, encoding, fastpath=True)
  File "~/.conda/envs/plotting/lib/python3.10/site-packages/xarray/core/variable.py", line 2668, in __init__
    raise ValueError(f"{type(self).__name__} objects must be 1-dimensional")
ValueError: IndexVariable objects must be 1-dimensional

Anything else we need to know?

No response

Environment

INSTALLED VERSIONS

commit: None
python: 3.10.4 | packaged by conda-forge | (main, Mar 24 2022, 17:39:04) [GCC 10.3.0]
python-bits: 64
OS: Linux
OS-release: 3.10.0-1160.59.1.el7.x86_64
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: 2022.3.0
pandas: 1.4.2
numpy: 1.22.3
scipy: 1.8.0
netCDF4: 1.5.8
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: None
cftime: 1.6.0
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: None
distributed: None
matplotlib: 3.5.1
cartopy: 0.20.2
seaborn: None
numbagg: None
fsspec: None
cupy: None
pint: None
sparse: None
setuptools: 61.3.1
pip: 22.0.4
conda: None
pytest: None
IPython: None
sphinx: None

[benbovy]

I tested your example with Xarray's main branch and it seems to work (at least there's no ValueError anymore). Maybe some side effect of the indexes refactor in #5692? I just run the example without inspecting the outputs so not sure if everything is right, though.

[headtr1ck]

Just a tip: You don't need any stacking for that.
Just use an indexer with a new dim:

diag = xr.DataArray(np.arange(12), dims="diag")
test.isel(dim0=diag, adj_dim0=diag)

this should do the trick ;)

[DWesl]

Just a tip: You don't need any stacking for that. Just use an indexer with a new dim:

I am aware that I can extract the diagonal of the arrays by using the same index for each argument of isel. That is, in fact, how I extracted the diagonals in each case above (look for diag_index to find the examples).

The bit that interests me is unstacking the relevant dimension, because the data in the original case comes to me with, effectively, a stacked dimension, and I would like to turn it back into an unstacked dimension because that is what I am used to using pcolormesh to plot.

That is to say, skipping the unstacking rather defeats the purpose of what I am trying to do, unless you have suggestions for how to create a two-dimensional plot (one using something like contourf or pcolormesh) of a one-dimensional Dataset, or a series of two-dimensional plots from a two-dimensional Dataset.

[benbovy]

@DWesl have you tried running your code with Xarray's main branch? If yes, could you confirm whether or not it yields the expected output? Thanks!

[DWesl]

Oh, right, you suggested that a bit ago.

When I checkout upstream/main in my local XArray repository root and run the example, it completes without error. When I fix the example to use the correct dimension, the implicit print on the last line shows the nearly the same as unstacked_diag from a few lines earlier.

Still not sure what fixed this, but since it's working, I don't care so much. I will wait for this to show up in a release. Thank you!