Repository: pydata/xarray

Issue:
Assigning values to a subset of a dataset
Hi, can somebody tell me what is the "correct" way to manipulate a subset of the data contained in a Dataset? 

Consider the following example:
```
import numpy as np
import xarray as xr

shape = (3, 2)
da1 = xr.DataArray(np.zeros(shape), dims=('x', 'y'), coords=dict(x=[1, 2, 3], y=[4, 5]), name='var1')
da2 = xr.DataArray(np.zeros(shape), dims=('x', 'y'), coords=dict(x=[1, 2, 3], y=[4, 5]), name='var2')
```

I can easily change the value of variable 1 at a given coordinate in the first DataArray using the following syntax:
```
da1.loc[dict(x=1, y=4)] = 1
```

However, if I merge both DataArrays into a single Dataset and want to change both variables at the same time, there seems to be no straightforward solution:
```
ds = xr.merge([da1, da2])
ds.loc[dict(x=1, y=4)] =  ... <-- what to write here?
```

The only solution I could come up with is to modify the two values separately, but this is neither very elegant nor scales with the number of variables:
```
ds['var1'].loc[dict(x=1, y=4)] = 2
ds['var2'].loc[dict(x=1, y=4)] = 3
```

All I could find in the docs about this issue is:
> Using indexing to assign values to a subset of dataset (e.g., ds[dict(space=0)] = 1) is not yet supported.

If not by indexing, what other (more compact) way exists? A potential solution might be to create a separate Dataset and then use the `update` method, but this seems overly complicated, too.



Discussion:
matzegoebel: I haven't updated the documentation yet, where it says that this feature is not supported yet. Do you think we need example code for this feature in the documentation?
max-sixty: @matzegoebel forgive the very long delay on the review. We're planning to find a better system to ensure these don't drop through.

I would be up for adding this, for consistency. I don't think I've ever needed the functionality, but it also doesn't make the interface more complicated given it's mirroring `__getitem__`.

We probably need to think through whether there are any corner cases here; I can't think of any atm.

Any other thoughts?
shoyer: Strong +1 from me on narrowing scope whenever possible. Adding features
incrementally is much easier than doing things all at once :)

On Sun, Apr 18, 2021 at 6:15 PM Maximilian Roos ***@***.***>
wrote:

> ***@***.**** commented on this pull request.
> ------------------------------
>
> In xarray/core/dataset.py
> <https://github.com/pydata/xarray/pull/5045#discussion_r615490870>:
>
> > +            # loop over dataset variables
> +            for var, da in self.items():
> +                val = value
> +                if type(value) == xr.core.dataset.Dataset:
> +                    val = value[var]
> +                # only set value if all dimensions are present
> +                if all([k in da.dims for k in key.keys()]):
> +                    da[key] = val
>
> V good point. Thanks.
>
> Either a) raising on missing dimensions or b) "entirely set to zero" in
> this case, would be reasonable imo.
>
> To the extent you're less confident that (b) is correct, I'd suggest we
> move forward with (a) and evaluating whether we should do (b) separately.
>
> (though @shoyer <https://github.com/shoyer> if you're up for it, let's
> discuss the general pace of reviews next week and whether you think this
> sort of "narrowing" of PR scope is a reasonable tactic)
>
> —
> You are receiving this because you were mentioned.
>
>
> Reply to this email directly, view it on GitHub
> <https://github.com/pydata/xarray/pull/5045#discussion_r615490870>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAJJFVQR6Q6VI44EEY2SF43TJN7ZTANCNFSM4ZKTDBUA>
> .
>

matzegoebel: ok I tried to include your suggestions. Concerning @shoyer's point 3: Since I guess there could be a lot of different errors appearing, I created a copy of the data to be changed to check if the setitem fails before doing the actual update. That's of course suboptimal concerning the performance. What do you think, should we include checks for all conceivable errors or keep this test update?
matzegoebel: I don't understand what's the issue of the failing test. Do you?

keewis: that's a flaky test which randomly fails (see #4539). You can safely ignore it, the CI should pass on the next run.
max-sixty: Great, this is shaping up.

I think we can find a way of failing early on bad indexes without attempting the whole operation on a copy.

At the very least, we could call `__getitem__` with the indexes and see whether that passes. There may be better ways yet.

I also think that because the currently proposed code uses a shallow copy, it may be mutating the original when bad indexes are passed — it's worth adding a test to confirm.
matzegoebel: Calling  __getitem__ is not enough to detect all possible errors, I guess. Another possibility would be to do a deep copy before the assignments, and if anything goes wrong, restore the original data from the copy. In this way, the assignments do not have to be done twice unless an error appears.
max-sixty: Which errors would `__getitem__` miss? Type errors that don't coerce.

The issue with a deep copy of the whole dataset is that it's very expensive. It's probably better to have that rather than nothing, but it could have confusing performance effects given that people are often going to be mutating values to reduce copies.

These aren't strongly held views though. Any thoughts from others? 
matzegoebel: ok, I deleted the copy stuff and included a few checks to catch possible errors before setting the values. Did I miss anything? How do we check for "type errors that don't coerce", as you mentioned? 
The setitem method of the LocIndexer now calls the setitem method of the Dataset class, so that we don't have redundant code.
max-sixty: > ok, I deleted the copy stuff and included a few checks to catch possible errors before setting the values. Did I miss anything? How do we check for "type errors that don't coerce", as you mentioned?

Excellent. Re the checks — I mostly meant that it was going to be very rare for something to get through — I don't think it's necessary to check for something like "type errors that don't coerce". 
matzegoebel: @shoyer thanks for your suggestions! I included them as best as I could.
matzegoebel: I revised the pre-assignment checks. In my opinion xr.align is not so helpful when checking that the dimension sizes and coordinates are consistent, because it doesn't fail when the dimension size of the two Datasets is different, but the coordinate of the second Dataset is a subset of the first one. Therefore, I reimplemented the check that I had previously in a similar way. I also added a check for the wrong order of the dimensions, that you mentioned @shoyer.
If, despite the checks, an error occurs during the assignment, e.g. due to a type error, and the dataset has been updated already partially, the user is informed about this.
shoyer: > I revised the pre-assignment checks. In my opinion xr.align is not so helpful when checking that the dimension sizes and coordinates are consistent, because it doesn't fail when the dimension size of the two Datasets is different, but the coordinate of the second Dataset is a subset of the first one.

Could you kindly elaborate on this issue, maybe with a specific example?

> If, despite the checks, an error occurs during the assignment, e.g. due to a type error, and the dataset has been updated already partially, the user is informed about this.

`np.can_cast` with `casting='unsafe` can check this. It sounds like this would probably be something good to add to our checks :)
matzegoebel: > Could you kindly elaborate on this issue, maybe with a specific example?
I think I somehow forgot the join="exact" when testing the functionality of xr.align. So nevermind, I'll reimplement again. :P

> `np.can_cast` with `casting='unsafe` can check this. It sounds like this would probably be something good to add to our checks :)
Ok good point. I'll give it a try.

matzegoebel: > `np.can_cast` with `casting='unsafe` can check this. It sounds like this would probably be something good to add to our checks :)

I'm not sure how to use this, because `np.can_cast` with unsafe casting mostly returns True, e.g. `np.can_cast(str, float, casting="unsafe")==True`. The error that is implemented in the tests would not be caught then. But I could explicitly try to cast the new values to the datatype of the original values with `.astype`. Would that be an option?

shoyer: > I'm not sure how to use this, because `np.can_cast` with unsafe casting mostly returns True, e.g. `np.can_cast(str, float, casting="unsafe")==True`. The error that is implemented in the tests would not be caught then.

Oh, good point, thanks for checking.

> But I could explicitly try to cast the new values to the datatype of the original values with `.astype`. Would that be an option?

Yes, I like this idea!
matzegoebel: ok done

keewis: could you also resolve the merge conflicts? It seems we can't "approve and run" the CI with conflicts.
matzegoebel: > could you also resolve the merge conflicts? It seems we can't "approve and run" the CI with conflicts.
Ok I resolved them
dcherian: I've fixed the location of the whats-new note. Is this ready to go in?
max-sixty: > I've fixed the location of the whats-new note. Is this ready to go in?

Yes. I have one typing question but we can merge regardless if needed

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 02287bf9..9ba9a7e9 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -22,10 +22,11 @@ v0.18.3 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Allow assigning values to a subset of a dataset using positional or label-based
+  indexing (:issue:`3015`, :pull:`5362`). By `Matthias Göbel <https://github.com/matzegoebel>`_.
 - Attempting to reduce a weighted object over missing dimensions now raises an error (:pull:`5362`).
   By `Mattia Almansi <https://github.com/malmans2>`_.
 
-
 Breaking changes
 ~~~~~~~~~~~~~~~~
 
@@ -141,6 +142,7 @@ Xianxiang Li, Zeb Nicholls, crusaderky, dschwoerer, johnomotani, keewis
 
 New Features
 ~~~~~~~~~~~~
+
 - apply ``combine_attrs`` on data variables and coordinate variables when concatenating
   and merging datasets and dataarrays (:pull:`4902`).
   By `Justus Magin <https://github.com/keewis>`_.
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index 85213567..1697a7c6 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -561,6 +561,17 @@ class _LocIndexer:
             raise TypeError("can only lookup dictionaries from Dataset.loc")
         return self.dataset.sel(key)
 
+    def __setitem__(self, key, value) -> None:
+        if not utils.is_dict_like(key):
+            raise TypeError(
+                "can only set locations defined by dictionaries from Dataset.loc."
+                f" Got: {key}"
+            )
+
+        # set new values
+        pos_indexers, _ = remap_label_indexers(self.dataset, key)
+        self.dataset[pos_indexers] = value
+
 
 class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
     """A multi-dimensional, in memory, array database.
@@ -1476,19 +1487,16 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
         else:
             return self._copy_listed(np.asarray(key))
 
-    @overload
-    def __setitem__(self, key: List[Hashable], value) -> None:
-        ...
-
-    @overload
-    def __setitem__(self, key: Hashable, value) -> None:
-        ...
-
-    def __setitem__(self, key, value) -> None:
+    def __setitem__(self, key: Union[Hashable, List[Hashable], Mapping], value) -> None:
         """Add an array to this dataset.
         Multiple arrays can be added at the same time, in which case each of
         the following operations is applied to the respective value.
 
+        If key is a dictionary, update all variables in the dataset
+        one by one with the given value at the given location.
+        If the given value is also a dataset, select corresponding variables
+        in the given value and in the dataset to be changed.
+
         If value is a `DataArray`, call its `select_vars()` method, rename it
         to `key` and merge the contents of the resulting dataset into this
         dataset.
@@ -1498,11 +1506,25 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
         variable.
         """
         if utils.is_dict_like(key):
-            raise NotImplementedError(
-                "cannot yet use a dictionary as a key to set Dataset values"
-            )
+            # check for consistency and convert value to dataset
+            value = self._setitem_check(key, value)
+            # loop over dataset variables and set new values
+            processed = []
+            for name, var in self.items():
+                try:
+                    var[key] = value[name]
+                    processed.append(name)
+                except Exception as e:
+                    if processed:
+                        raise RuntimeError(
+                            "An error occured while setting values of the"
+                            f" variable '{name}'. The following variables have"
+                            f" been successfully updated:\n{processed}"
+                        ) from e
+                    else:
+                        raise e
 
-        if isinstance(key, list):
+        elif isinstance(key, list):
             if len(key) == 0:
                 raise ValueError("Empty list of variables to be set")
             if len(key) == 1:
@@ -1524,6 +1546,69 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
         else:
             self.update({key: value})
 
+    def _setitem_check(self, key, value):
+        """Consistency check for __setitem__
+
+        When assigning values to a subset of a Dataset, do consistency check beforehand
+        to avoid leaving the dataset in a partially updated state when an error occurs.
+        """
+        from .dataarray import DataArray
+
+        if isinstance(value, Dataset):
+            missing_vars = [
+                name for name in value.data_vars if name not in self.data_vars
+            ]
+            if missing_vars:
+                raise ValueError(
+                    f"Variables {missing_vars} in new values"
+                    f" not available in original dataset:\n{self}"
+                )
+        elif not any([isinstance(value, t) for t in [DataArray, Number, str]]):
+            raise TypeError(
+                "Dataset assignment only accepts DataArrays, Datasets, and scalars."
+            )
+
+        new_value = xr.Dataset()
+        for name, var in self.items():
+            # test indexing
+            try:
+                var_k = var[key]
+            except Exception as e:
+                raise ValueError(
+                    f"Variable '{name}': indexer {key} not available"
+                ) from e
+
+            if isinstance(value, Dataset):
+                val = value[name]
+            else:
+                val = value
+
+            if isinstance(val, DataArray):
+                # check consistency of dimensions
+                for dim in val.dims:
+                    if dim not in var_k.dims:
+                        raise KeyError(
+                            f"Variable '{name}': dimension '{dim}' appears in new values "
+                            f"but not in the indexed original data"
+                        )
+                dims = tuple([dim for dim in var_k.dims if dim in val.dims])
+                if dims != val.dims:
+                    raise ValueError(
+                        f"Variable '{name}': dimension order differs between"
+                        f" original and new data:\n{dims}\nvs.\n{val.dims}"
+                    )
+            else:
+                val = np.array(val)
+
+            # type conversion
+            new_value[name] = val.astype(var_k.dtype, copy=False)
+
+        # check consistency of dimension sizes and dimension coordinates
+        if isinstance(value, DataArray) or isinstance(value, Dataset):
+            xr.align(self[key], value, join="exact", copy=False)
+
+        return new_value
+
     def __delitem__(self, key: Hashable) -> None:
         """Remove a variable from this dataset."""
         del self._variables[key]


Test patch:
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 8d4f5d04..1c32366e 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -1721,8 +1721,6 @@ class TestDataset:
         assert_identical(expected, actual)
         with pytest.raises(TypeError, match=r"can only lookup dict"):
             data.loc["a"]
-        with pytest.raises(TypeError):
-            data.loc[dict(dim3="a")] = 0
 
     def test_selection_multiindex(self):
         mindex = pd.MultiIndex.from_product(
@@ -3422,8 +3420,60 @@ class TestDataset:
         data1["A"] = 3 * data2["A"]
         assert_equal(data1["A"], 3 * data2["A"])
 
-        with pytest.raises(NotImplementedError):
-            data1[{"x": 0}] = 0
+        # test assignment with positional and label-based indexing
+        data3 = data1[["var1", "var2"]]
+        data3["var3"] = data3.var1.isel(dim1=0)
+        data4 = data3.copy()
+        err_msg = (
+            "can only set locations defined by dictionaries from Dataset.loc. Got: a"
+        )
+        with pytest.raises(TypeError, match=err_msg):
+            data1.loc["a"] = 0
+        err_msg = r"Variables \['A', 'B', 'scalar'\] in new values not available in original dataset:"
+        with pytest.raises(ValueError, match=err_msg):
+            data4[{"dim2": 1}] = data1[{"dim2": 2}]
+        err_msg = "Variable 'var3': indexer {'dim2': 0} not available"
+        with pytest.raises(ValueError, match=err_msg):
+            data1[{"dim2": 0}] = 0.0
+        err_msg = "Variable 'var1': indexer {'dim2': 10} not available"
+        with pytest.raises(ValueError, match=err_msg):
+            data4[{"dim2": 10}] = data3[{"dim2": 2}]
+        err_msg = "Variable 'var1': dimension 'dim2' appears in new values"
+        with pytest.raises(KeyError, match=err_msg):
+            data4[{"dim2": 2}] = data3[{"dim2": [2]}]
+        err_msg = (
+            "Variable 'var2': dimension order differs between original and new data"
+        )
+        data3["var2"] = data3["var2"].T
+        with pytest.raises(ValueError, match=err_msg):
+            data4[{"dim2": [2, 3]}] = data3[{"dim2": [2, 3]}]
+        data3["var2"] = data3["var2"].T
+        err_msg = "indexes along dimension 'dim2' are not equal"
+        with pytest.raises(ValueError, match=err_msg):
+            data4[{"dim2": [2, 3]}] = data3[{"dim2": [2, 3, 4]}]
+        err_msg = "Dataset assignment only accepts DataArrays, Datasets, and scalars."
+        with pytest.raises(TypeError, match=err_msg):
+            data4[{"dim2": [2, 3]}] = data3["var1"][{"dim2": [3, 4]}].values
+        data5 = data4.astype(str)
+        data5["var4"] = data4["var1"]
+        err_msg = "could not convert string to float: 'a'"
+        with pytest.raises(ValueError, match=err_msg):
+            data5[{"dim2": 1}] = "a"
+
+        data4[{"dim2": 0}] = 0.0
+        data4[{"dim2": 1}] = data3[{"dim2": 2}]
+        data4.loc[{"dim2": 1.5}] = 1.0
+        data4.loc[{"dim2": 2.0}] = data3.loc[{"dim2": 2.5}]
+        for v, dat3 in data3.items():
+            dat4 = data4[v]
+            assert_array_equal(dat4[{"dim2": 0}], 0.0)
+            assert_array_equal(dat4[{"dim2": 1}], dat3[{"dim2": 2}])
+            assert_array_equal(dat4.loc[{"dim2": 1.5}], 1.0)
+            assert_array_equal(dat4.loc[{"dim2": 2.0}], dat3.loc[{"dim2": 2.5}])
+            unchanged = [1.0, 2.5, 3.0, 3.5, 4.0]
+            assert_identical(
+                dat4.loc[{"dim2": unchanged}], dat3.loc[{"dim2": unchanged}]
+            )
 
     def test_setitem_pandas(self):