Assigning with .loc indexing: implicit effect of dimension order?

[darikg]

What happened?

Assigning a DataArray to a subset of another DataArray seems to depend on the order of the dimensions

What did you expect to happen?

I would expect xarray to automatically align the dimensions appropriately.

Minimal Complete Verifiable Example

Example 1:

from xarray import DataArray
from numpy import zeros, ones, arange

# Two arrays abc and acb, identical except for dimensional ordering
abc = DataArray(zeros((2, 3, 4)), dims=('a', 'b', 'c'))
acb = DataArray(zeros((2, 4, 3)), dims=('a', 'c', 'b'))
assert (abc == acb).all()

# Assign a subset
bc = DataArray(ones((3, 4)), dims=('b', 'c'))
abc.loc[dict(a=0)] = bc
acb.loc[dict(a=0)] = bc # ValueError: could not broadcast input array from shape (3,4) into shape (4,3)

Example 2:

Same as example 1, but instead of raising a ValueError, runs without error, but gives incorrect results because it ignores dimension order

# This time, make b and c dimensions identical
abc = DataArray(zeros((2, 3, 3)), dims=('a', 'b', 'c'))
acb = DataArray(zeros((2, 3, 3)), dims=('a', 'c', 'b'))
assert (abc == acb).all()

# Assign a subset
bc = DataArray(arange(9).reshape(3, 3), dims=('b', 'c'))
abc.loc[dict(a=0)] = bc
acb.loc[dict(a=0)] = bc
assert (abc == acb).all()  # Assertion error

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.

Relevant log output

No response

Anything else we need to know?

No response

Environment

INSTALLED VERSIONS ------------------ commit: None python: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:50:36) [MSC v.1929 64 bit (AMD64)] python-bits: 64 OS: Windows OS-release: 10 machine: AMD64 processor: Intel64 Family 6 Model 142 Stepping 12, GenuineIntel byteorder: little LC_ALL: None LANG: en_US.UTF-8 LOCALE: ('English_United States', '1252') libhdf5: 1.12.1 libnetcdf: 4.8.1

xarray: 2022.6.0
pandas: 1.4.3
numpy: 1.23.1
scipy: 1.9.0
netCDF4: 1.6.0
pydap: None
h5netcdf: None
h5py: 3.7.0
Nio: None
zarr: None
cftime: 1.6.1
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: None
distributed: None
matplotlib: 3.5.2
cartopy: None
seaborn: None
numbagg: None
fsspec: 2022.7.1
cupy: None
pint: None
sparse: None
flox: None
numpy_groupies: None
setuptools: 63.4.2
pip: 22.2.2
conda: 4.14.0
pytest: 7.1.2
IPython: 8.4.0
sphinx: None

[maarten-vandenberg]

We are experiencing the same issue (with xarray version 2022.3.0). After a bit of exploring it seems that in the Variable __setitem__ method, dimensions of the value get switched around. This leads to the behaviour observed in the examples above.

[dcherian]

Thanks @maarten-vandenberg . It looks like you have made an important step toward solving the problem.

Do you have time to open a pull request fixing this issue? Even an incomplete one helps get us closer to a fix.

[maarten-vandenberg]

I'd be happy to open a pull request, although I must admit I'm not very familiar with the internals of xarray. I hope it will help to clarify the issue and bring a potential solution closer

[dranjan]

I also ran into this. Here's a workaround that seems to work in some cases, which uses a Dataset and explicit broadcasting:

index = xarray.Dataset({'a': 0})
index, value = xarray.broadcast(index, bc)
acb.loc[index] = value
assert (abc == acb).all()

This doesn't work in all cases, such as when acb has additional dimensions that are not being sliced. I expect broadcasting against the indexed result would work more reliably, but it seems pretty wasteful because it has to perform the indexing explicitly:

value = bc.broadcast_like(acb[dict(a=0)])
acb[dict(a=0)] = value
assert (abc == acb).all()

[dranjan]

I attempted to solve this in #8067. Until/unless that gets accepted, I think an acceptable workaround could be

acb.loc[dict(a=0)] = bc.variable