Repository: pydata/xarray

Issue:
Test test_resample_loffset fails on i586 with Python 3.9
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports: http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples: https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
Running tests with `pytest-3.9 -v -n auto xarray`
 ->
```
[  262s] =================================== FAILURES ===================================
[  262s] ______________________ TestDataset.test_resample_loffset _______________________
[  262s] [gw1] linux -- Python 3.9.5 /usr/bin/python3.9
[  262s] 
[  262s] self = <xarray.tests.test_dataset.TestDataset object at 0xf0719a18>
[  262s] 
[  262s]     def test_resample_loffset(self):
[  262s]         times = pd.date_range("2000-01-01", freq="6H", periods=10)
[  262s]         ds = Dataset(
[  262s]             {
[  262s]                 "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)),
[  262s]                 "bar": ("time", np.random.randn(10), {"meta": "data"}),
[  262s]                 "time": times,
[  262s]             }
[  262s]         )
[  262s]         ds.attrs["dsmeta"] = "dsdata"
[  262s]     
[  262s]         # Our use of `loffset` may change if we align our API with pandas' changes.
[  262s]         # ref https://github.com/pydata/xarray/pull/4537
[  262s]         actual = ds.resample(time="24H", loffset="-12H").mean().bar
[  262s]         expected_ = ds.bar.to_series().resample("24H").mean()
[  262s]         expected_.index += to_offset("-12H")
[  262s]         expected = DataArray.from_series(expected_)
[  262s] >       assert_identical(actual, expected)
[  262s] E       AssertionError: Left and right DataArray objects are not identical
[  262s] E       
[  262s] E       Differing values:
[  262s] E       L
[  262s] E           array([ 0.018477, -0.178807, -0.792644])
[  262s] E       R
[  262s] E           array([ 0.018477, -0.178807, -0.792644])
[  262s] 
[  262s] /home/abuild/rpmbuild/BUILD/xarray-0.18.1/xarray/tests/test_dataset.py:3991: AssertionError
```
[devel:languages:python:numeric-python-xarray-openSUSE_Tumbleweed-i586.txt](Complete log)

**What you expected to happen**:
Tests succeed

**Minimal Complete Verifiable Example**:

```python
# Put your MCVE code here
```

**Anything else we need to know?**:
Tests with x86_64 succeed

**Environment**:

openSUSE Tumbleweed build environment (for packaging)

<details><summary>Output of <tt>xr.show_versions()</tt></summary>

```

INSTALLED VERSIONS
------------------
commit: None
python: 3.9.5 (default, May 05 2021, 15:23:13) [GCC]
python-bits: 32
OS: Linux
OS-release: 5.12.3-1-default
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: None
LOCALE: ('en_US', 'UTF-8')
libhdf5: None
libnetcdf: None

xarray: 0.18.1
pandas: 1.2.4
numpy: 1.20.2
scipy: 1.6.3
netCDF4: None
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: None
cftime: None
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: 2021.04.1
distributed: None
matplotlib: None
cartopy: None
seaborn: None
numbagg: None
pint: None
setuptools: 44.1.1
pip: None
conda: None
pytest: 6.2.2
IPython: None
sphinx: None
```

</details>


Discussion:
max-sixty: Thanks @sebix 

Do we know where the difference is coming from? If I'm reading the test correctly, the floats aren't related to the offset — which is on discrete values — and so this may be a broader problem. Is that right?
sebix: > Thanks @sebix
> 
> Do we know where the difference is coming from? If I'm reading the test correctly, the floats aren't related to the offset — which is on discrete values — and so this may be a broader problem. Is that right?

Sorry, I have no more insight, I just know that the test is failing on that particular architecture (x86_64, ppc64, ppc64le, aarch64 are fine).

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 828e8da4..9ba9a7e9 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -22,6 +22,10 @@ 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
 ~~~~~~~~~~~~~~~~
@@ -138,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]
diff --git a/xarray/core/formatting.py b/xarray/core/formatting.py
index c3e7187a..a3c21bdd 100644
--- a/xarray/core/formatting.py
+++ b/xarray/core/formatting.py
@@ -258,12 +258,12 @@ def inline_variable_array_repr(var, max_width):
     """Build a one-line summary of a variable's data."""
     if var._in_memory:
         return format_array_flat(var, max_width)
+    elif hasattr(var._data, "_repr_inline_"):
+        return var._data._repr_inline_(max_width)
     elif isinstance(var._data, dask_array_type):
         return inline_dask_repr(var.data)
     elif isinstance(var._data, sparse_array_type):
         return inline_sparse_repr(var.data)
-    elif hasattr(var._data, "_repr_inline_"):
-        return var._data._repr_inline_(max_width)
     elif hasattr(var._data, "__array_function__"):
         return maybe_truncate(repr(var._data).replace("\n", " "), max_width)
     else:
diff --git a/xarray/core/weighted.py b/xarray/core/weighted.py
index a3ba0396..e8838b07 100644
--- a/xarray/core/weighted.py
+++ b/xarray/core/weighted.py
@@ -119,6 +119,19 @@ class Weighted(Generic[T_DataWithCoords]):
         self.obj: T_DataWithCoords = obj
         self.weights: "DataArray" = weights
 
+    def _check_dim(self, dim: Optional[Union[Hashable, Iterable[Hashable]]]):
+        """raise an error if any dimension is missing"""
+
+        if isinstance(dim, str) or not isinstance(dim, Iterable):
+            dims = [dim] if dim else []
+        else:
+            dims = list(dim)
+        missing_dims = set(dims) - set(self.obj.dims) - set(self.weights.dims)
+        if missing_dims:
+            raise ValueError(
+                f"{self.__class__.__name__} does not contain the dimensions: {missing_dims}"
+            )
+
     @staticmethod
     def _reduce(
         da: "DataArray",
@@ -236,6 +249,8 @@ class Weighted(Generic[T_DataWithCoords]):
 class DataArrayWeighted(Weighted["DataArray"]):
     def _implementation(self, func, dim, **kwargs) -> "DataArray":
 
+        self._check_dim(dim)
+
         dataset = self.obj._to_temp_dataset()
         dataset = dataset.map(func, dim=dim, **kwargs)
         return self.obj._from_temp_dataset(dataset)
@@ -244,6 +259,8 @@ class DataArrayWeighted(Weighted["DataArray"]):
 class DatasetWeighted(Weighted["Dataset"]):
     def _implementation(self, func, dim, **kwargs) -> "Dataset":
 
+        self._check_dim(dim)
+
         return self.obj.map(func, dim=dim, **kwargs)
 
 


Test patch:
diff --git a/xarray/tests/test_cupy.py b/xarray/tests/test_cupy.py
index 0276b8eb..69f43d99 100644
--- a/xarray/tests/test_cupy.py
+++ b/xarray/tests/test_cupy.py
@@ -47,7 +47,7 @@ def test_cupy_import():
 def test_check_data_stays_on_gpu(toy_weather_data):
     """Perform some operations and check the data stays on the GPU."""
     freeze = (toy_weather_data["tmin"] <= 0).groupby("time.month").mean("time")
-    assert isinstance(freeze.data, cp.core.core.ndarray)
+    assert isinstance(freeze.data, cp.ndarray)
 
 
 def test_where():
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 8d4f5d04..00a6fb82 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):
 
@@ -3987,7 +4037,7 @@ class TestDataset:
         expected_ = ds.bar.to_series().resample("24H").mean()
         expected_.index += to_offset("-12H")
         expected = DataArray.from_series(expected_)
-        assert_identical(actual, expected)
+        assert_allclose(actual, expected)
 
     def test_resample_by_mean_discarding_attrs(self):
         times = pd.date_range("2000-01-01", freq="6H", periods=10)
diff --git a/xarray/tests/test_weighted.py b/xarray/tests/test_weighted.py
index 31c3c552..45e662f1 100644
--- a/xarray/tests/test_weighted.py
+++ b/xarray/tests/test_weighted.py
@@ -368,3 +368,19 @@ def test_weighted_operations_keep_attr_da_in_ds(operation):
     result = getattr(data.weighted(weights), operation)(keep_attrs=True)
 
     assert data.a.attrs == result.a.attrs
+
+
+@pytest.mark.parametrize("as_dataset", (True, False))
+def test_weighted_bad_dim(as_dataset):
+
+    data = DataArray(np.random.randn(2, 2))
+    weights = xr.ones_like(data)
+    if as_dataset:
+        data = data.to_dataset(name="data")
+
+    error_msg = (
+        f"{data.__class__.__name__}Weighted"
+        " does not contain the dimensions: {'bad_dim'}"
+    )
+    with pytest.raises(ValueError, match=error_msg):
+        data.weighted(weights).mean("bad_dim")