Setting item with loc and boolean mask fails

[niowniow]

What happened:
When setting a DataArray with loc, xarray converts bool masks to the coord data type. Therefore loc does not use a boolean mask but tries to match the indexes of the coord.

Minimal Complete Verifiable Example:

import numpy as np
import xarray as xr


x = np.arange(10).astype(np.float64)
fx = np.arange(10).astype(np.float64)
da = xr.DataArray(fx,dims=['x'],coords={'x':x})
mask = np.zeros((10,))
mask[1::2] = 1
mask = mask.astype(bool)
da.loc[{'x':~mask}] = np.arange(5)+10

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/matthmey/miniconda3/envs/analytics/lib/python3.7/site-packages/xarray/core/dataarray.py", line 214, in __setitem__
    self.data_array[pos_indexers] = value
  File "/home/matthmey/miniconda3/envs/analytics/lib/python3.7/site-packages/xarray/core/dataarray.py", line 767, in __setitem__
    self.variable[key] = value
  File "/home/matthmey/miniconda3/envs/analytics/lib/python3.7/site-packages/xarray/core/variable.py", line 854, in __setitem__
    indexable[index_tuple] = value
  File "/home/matthmey/miniconda3/envs/analytics/lib/python3.7/site-packages/xarray/core/indexing.py", line 1164, in __setitem__
    array[key] = value
ValueError: shape mismatch: value array of shape (5,) could not be broadcast to indexing result of shape (10,)

Anything else we need to know?:
Could be fixed by replacing the line, but maybe this is not the cleanest solution. I tried fixing it with the following, which works for the above code but fails for other cases.

            if label.dtype != np.bool:
                label = maybe_cast_to_coords_dtype(label, coord.dtype)

Environment:

Output of xr.show_versions()

INSTALLED VERSIONS

commit: None
python: 3.7.0 (default, Oct 9 2018, 10:31:47)
[GCC 7.3.0]
python-bits: 64
OS: Linux
OS-release: 5.4.0-80-lowlatency
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: None
libnetcdf: None

xarray: 0.19.0
pandas: 1.0.1
numpy: 1.18.1
scipy: 1.4.1
netCDF4: None
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: 2.5.0
cftime: None
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: 2.14.0
distributed: None
matplotlib: 3.1.3
cartopy: None
seaborn: None
numbagg: None
pint: None
setuptools: 46.1.3.post20200330
pip: 20.0.2
conda: None
pytest: None
IPython: 7.13.0
sphinx: None

[benbovy]

With da.loc[{'x':~mask}], you're actually doing label-based indexing, so ~mask will select repeatedly the items for which x=1.0 (True) and x=0.0 (False):

>>> da.loc[{'x':~mask}]
<xarray.DataArray (x: 10)>
array([1., 0., 1., 0., 1., 0., 1., 0., 1., 0.])
Coordinates:
  * x        (x) float64 1.0 0.0 1.0 0.0 1.0 0.0 1.0 0.0 1.0 0.0

What you want is integer-based indexing (i.e., not using .loc):

>>> da[{'x':~mask}]
<xarray.DataArray (x: 5)>
array([0., 2., 4., 6., 8.])
Coordinates:
  * x        (x) float64 0.0 2.0 4.0 6.0 8.0

doc link.

[niowniow]

Thank you! You are right. I guess based on the docs for Assigning values with indexing I assumed only .loc features item assignment . I vaguely remember there where issues with assigning via integer-based indexing. Well, propably my bad.

Should we add an example how to perform masking this way to the docs? There are many examples about masking using different ways, but not this simple, intuitive way.

Even now, it's not really clear to me why "boolean"-masking should only work with integer-indexing and not label-indexing. Given that the highlighted maybe_cast_to_coords_dtype is the only reason it does not work for label-based indexing.

edit: probably the reason is to support label-based indexing for boolean coords?!

[benbovy]

I agree that this kind of "boolean indexing" could be clarified a bit in the docs.

I also had to double check here since I use isel() / sel() much more often than [] and .loc[].

Even now, it's not really clear to me why "boolean"-masking should only work with integer-indexing and not label-indexing. Given that the highlighted maybe_cast_to_coords_dtype is the only reason it does not work for label-based indexing.

Hmm you're right. The primary reason of maybe_cast_to_coords_dtype is because pandas indexes don't support floating point values with lower precision than float64 (e.g., #5700). But here

xarray/xarray/core/indexes.py

Lines 236 to 237 in a6b44d7

	elif label.dtype.kind == "b":
	indexer = label

we actually bypass the underlying pandas index when the labels are a boolean array-like so that boolean indexing should normally work with both integer and label based indexing (#182).

So we should probably fix maybe_cast_to_coords_dtype with something like

def maybe_cast_to_coords_dtype(label, coords_dtype):
    label_dtype = np.dtype(getattr(label, "dtype", None))
    if coords_dtype.kind == "f" and not isinstance(label, slice) and label_dtype.kind != "b":
        label = np.asarray(label, dtype=coords_dtype)
    return label