Equivalent to np.isclose()?

[ahuang11]

Sometimes when I do ds.where(ds['lat'] == 15, drop=True), it doesn't find any due to some issue with floating point equality check so I have to use np.isclose.

However, when I do this:

ds = ds.where(np.isclose(ds['lat'], 15), drop=True)

I get:

TypeError: cond argument is array([False, False, False, ..., False, False, False]) but must be a <class 'xarray.core.dataset.Dataset'> or <class 'xarray.core.dataarray.DataArray'>

So I have to do this:

        ds = ds.where(
            xr.DataArray(np.isclose(ds['lat'], 15), dims='index'),
            drop=True,
        )

[max-sixty]

Hi @ahuang11 , thanks for the question

Would .sel(tolerance={}) be a good solution? It requires specifying a number, rather than the easier approach of delegating that to numpy.

I also think we could have more support for indexing with boolean arrays, such that ds.where(np.isclose(ds['lat'], 15), drop=True) works. IIRC we support this in some indexing operations, though would have to check whether we support numpy arrays.

[ahuang11]

I guess it would work, but I would still have to swap_dims e.g. ds['lat'].swap_dims({'index': 'lat'}).sel(lat=15, tolerance=0.01).swap_dims({'lat': 'index'}) which is very verbose.

Support for indexing with boolean arrays would be nice. Where would this be implemented?

[max-sixty]

Ah, yes, good point — if lat isn't an index, then that verbose swap is required.

So a few issues:

• Boolean indexing: I had thought this was implemented, but below is a case of it not working — does anyone have a better memory than me on where / whether this works?
• Allowing a numpy array to ds.where
• Indexing on non-index coords — there's lots of discussion on this, and there's a wider project to redo how we do indexing

In [4]: ds = xr.tutorial.open_dataset("air_temperature", chunks={"time": -1}).isel(time=slice(100,200))

In [5]: ds
Out[5]:
<xarray.Dataset>
Dimensions:  (lat: 25, lon: 53, time: 100)
Coordinates:
  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0
  * time     (time) datetime64[ns] 2013-01-26 ... 2013-02-19T18:00:00
Data variables:
    air      (time, lat, lon) float32 dask.array<chunksize=(100, 25, 53), meta=np.ndarray>
Attributes:
    Conventions:  COARDS
    title:        4x daily NMC reanalysis (1948)
    description:  Data is from NMC initialized reanalysis\n(4x/day).  These a...
    platform:     Model
    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...

In [6]: ds.sel(lat=ds['lat']>50)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-6-cf01f91526ec> in <module>
----> 1 ds.sel(lat=ds['lat']>50)

~/Library/Caches/pypoetry/virtualenvs/-cyukv0jK-py3.8/lib/python3.8/site-packages/xarray/core/dataset.py in sel(self, indexers, method, tolerance, drop, **indexers_kwargs)
   2180         """
   2181         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel")
-> 2182         pos_indexers, new_indexes = remap_label_indexers(
   2183             self, indexers=indexers, method=method, tolerance=tolerance
   2184         )

~/Library/Caches/pypoetry/virtualenvs/-cyukv0jK-py3.8/lib/python3.8/site-packages/xarray/core/coordinates.py in remap_label_indexers(obj, indexers, method, tolerance, **indexers_kwargs)
    395     }
    396
--> 397     pos_indexers, new_indexes = indexing.remap_label_indexers(
    398         obj, v_indexers, method=method, tolerance=tolerance
    399     )

~/Library/Caches/pypoetry/virtualenvs/-cyukv0jK-py3.8/lib/python3.8/site-packages/xarray/core/indexing.py in remap_label_indexers(data_obj, indexers, method, tolerance)
    268             coords_dtype = data_obj.coords[dim].dtype
    269             label = maybe_cast_to_coords_dtype(label, coords_dtype)
--> 270             idxr, new_idx = convert_label_indexer(index, label, dim, method, tolerance)
    271             pos_indexers[dim] = idxr
    272             if new_idx is not None:

~/Library/Caches/pypoetry/virtualenvs/-cyukv0jK-py3.8/lib/python3.8/site-packages/xarray/core/indexing.py in convert_label_indexer(index, label, index_name, method, tolerance)
    200             indexer = get_indexer_nd(index, label, method, tolerance)
    201             if np.any(indexer < 0):
--> 202                 raise KeyError(f"not all values found in index {index_name!r}")
    203     return indexer, new_index
    204

KeyError: "not all values found in index 'lat'"

[keewis]

we could also fix #3917 and have np.isclose(xarray_object) return a xarray object, but this is probably quite a bit of work.