cf-xarray
cf-xarray copied to clipboard
Support dictionary-style indexing with `.loc`
Given the following sample data:
import numpy as np
import xarray as xr
import cf_xarray
lev = np.arange(1, 11)
lev = xr.DataArray(lev, dims='lev', attrs={'units': 'km', 'axis': 'Z'})
data = np.arange(lev.size)
data = xr.DataArray(data, dims='lev', coords={'lev': lev})
print(data.cf.axes) # prints {'Z': ['lev']}
print(data[{'lev': 0}]) # prints data for level value "1"
print(data.loc[{'lev': 1}]) # prints data for level value "1"
The cf accessor cannot handle dictionary-style position-based indexing:
>>> data.cf[{'Z': 0}]
File ~/miniconda3/lib/python3.10/site-packages/cf_xarray/accessor.py:2402, in CFDataArrayAccessor.__getitem__(self, key)
2370 """
2371 Index into a DataArray making use of CF attributes.
2372
(...)
2398 Add additional keys by specifying "custom criteria". See :ref:`custom_criteria` for more.
2399 """
2401 if not isinstance(key, str):
-> 2402 raise KeyError(
2403 f"Cannot use a list of keys with DataArrays. Expected a single string. Received {key!r} instead."
2404 )
2406 return _getitem(self, key)
KeyError: "Cannot use a list of keys with DataArrays. Expected a single string. Received {'Z': 0} instead."
It also cannot handle dictionary-style coordinate-based inexing:
>>> data.cf.loc[{'Z': 0}]
File ~/miniconda3/lib/python3.10/site-packages/cf_xarray/accessor.py:576, in _getattr(obj, attr, accessor, key_mappers, wrap_classes, extra_decorator)
573 func: Callable = attribute
575 else:
--> 576 raise AttributeError(
577 f"cf_xarray does not know how to wrap attribute '{type(obj).__name__}.{attr}'. "
578 "Please file an issue if you have a solution."
579 )
581 @functools.wraps(func)
582 def wrapper(*args, **kwargs):
583 posargs, arguments = accessor._process_signature(
584 func, args, kwargs, key_mappers
585 )
AttributeError: cf_xarray does not know how to wrap attribute 'DataArray.loc'. Please file an issue if you have a solution.
I'd be interested in addressing this but don't currently have time. Thought I'd just get the issue on this project's radar for now (if it wasn't already / I didn't miss an already-open thread).
:+1: should be straightforward. I don't use this syntax, that's why it hasn't been implemented yet :)