Skip to content

Commit

Permalink
Fix assignment with .loc (pydata#8067)
Browse files Browse the repository at this point in the history
* Add unit test for Variable[...] = DataArray (pydata#7030)

* Check for DataArray instance (pydata#7030)

* Update whats-new.rst (pydata#7030)

* Revert Variable and fix DataArray instead (pydata#7030)

* Add test case for .loc[...] = dataarray (pydata#7030)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix whats-new.rst (pydata#7030)

* Add regression test for Dataset (pydata#7030)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Maximilian Roos <[email protected]>
  • Loading branch information
3 people authored Sep 12, 2023
1 parent de66dae commit 8215911
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 0 deletions.
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,9 @@ Bug fixes
names were not updated properly internally (:issue:`7405`, :issue:`7588`,
:pull:`8104`).
By `Benoît Bovy <https://github.com/benbovy>`_.
- Fix bug where :py:class:`DataArray` instances on the right-hand side
of :py:meth:`DataArray.__setitem__` lose dimension names.
(:issue:`7030`, :pull:`8067`) By `Darsh Ranjan <https://github.com/dranjan>`_.

Documentation
~~~~~~~~~~~~~
Expand Down
1 change: 1 addition & 0 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -856,6 +856,7 @@ def __setitem__(self, key: Any, value: Any) -> None:
obj = self[key]
if isinstance(value, DataArray):
assert_coordinate_consistent(value, obj.coords.variables)
value = value.variable
# DataArray key -> Variable key
key = {
k: v.variable if isinstance(v, DataArray) else v
Expand Down
21 changes: 21 additions & 0 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -841,6 +841,27 @@ def get_data():
)
da[dict(x=ind)] = value # should not raise

def test_setitem_vectorized(self) -> None:
# Regression test for GH:7030
# Positional indexing
v = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])
b = xr.DataArray([[0, 0], [1, 0]], dims=["u", "v"])
c = xr.DataArray([[0, 1], [2, 3]], dims=["u", "v"])
w = xr.DataArray([-1, -2], dims=["u"])
index = dict(b=b, c=c)
v[index] = w
assert (v[index] == w).all()

# Indexing with coordinates
v = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])
v.coords["b"] = [2, 4, 6]
b = xr.DataArray([[2, 2], [4, 2]], dims=["u", "v"])
c = xr.DataArray([[0, 1], [2, 3]], dims=["u", "v"])
w = xr.DataArray([-1, -2], dims=["u"])
index = dict(b=b, c=c)
v.loc[index] = w
assert (v.loc[index] == w).all()

def test_contains(self) -> None:
data_array = DataArray([1, 2])
assert 1 in data_array
Expand Down
23 changes: 23 additions & 0 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4202,6 +4202,29 @@ def test_setitem_align_new_indexes(self) -> None:
)
assert_identical(ds, expected)

def test_setitem_vectorized(self) -> None:
# Regression test for GH:7030
# Positional indexing
da = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])
ds = xr.Dataset({"da": da})
b = xr.DataArray([[0, 0], [1, 0]], dims=["u", "v"])
c = xr.DataArray([[0, 1], [2, 3]], dims=["u", "v"])
w = xr.DataArray([-1, -2], dims=["u"])
index = dict(b=b, c=c)
ds[index] = xr.Dataset({"da": w})
assert (ds[index]["da"] == w).all()

# Indexing with coordinates
da = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])
ds = xr.Dataset({"da": da})
ds.coords["b"] = [2, 4, 6]
b = xr.DataArray([[2, 2], [4, 2]], dims=["u", "v"])
c = xr.DataArray([[0, 1], [2, 3]], dims=["u", "v"])
w = xr.DataArray([-1, -2], dims=["u"])
index = dict(b=b, c=c)
ds.loc[index] = xr.Dataset({"da": w}, coords={"b": ds.coords["b"]})
assert (ds.loc[index]["da"] == w).all()

@pytest.mark.parametrize("dtype", [str, bytes])
def test_setitem_str_dtype(self, dtype) -> None:
ds = xr.Dataset(coords={"x": np.array(["x", "y"], dtype=dtype)})
Expand Down

0 comments on commit 8215911

Please sign in to comment.