def test_either_dict_or_kwargs():
result = either_dict_or_kwargs(dict(a=1), None, 'foo')
expected = dict(a=1)
assert result == expected
result = either_dict_or_kwargs(None, dict(a=1), 'foo')
expected = dict(a=1)
assert result == expected
with pytest.raises(ValueError, match=r'foo'):
result = either_dict_or_kwargs(dict(a=1), dict(a=1), 'foo')
def test_either_dict_or_kwargs():
result = either_dict_or_kwargs(dict(a=1), None, "foo")
expected = dict(a=1)
assert result == expected
result = either_dict_or_kwargs(None, dict(a=1), "foo")
expected = dict(a=1)
assert result == expected
with pytest.raises(ValueError, match=r"foo"):
result = either_dict_or_kwargs(dict(a=1), dict(a=1), "foo")
def sel(self,
indexers=None,
method=None,
tolerance=None,
drop=False,
**indexers_kwargs) -> 'SparseDataArray':
"""Return a new array by selecting labels along the specified dim(s).
Overrides :meth:`~xarray.DataArray.sel` to handle >1-D indexers with
sparse data.
"""
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'sel')
if isinstance(indexers, dict) and len(indexers) > 1:
result = self
for k, v in indexers.items():
result = result.sel({k: v},
method=method,
tolerance=tolerance,
drop=drop)
return result
else:
return super().sel(indexers=indexers,
method=method,
tolerance=tolerance,
drop=drop)
def sel(self,
indexers=None,
method=None,
tolerance=None,
drop=False,
**indexers_kwargs):
"""Wrap Dataset.sel to handle units."""
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'sel')
indexers = _reassign_quantity_indexer(self._dataset, indexers)
return self._dataset.sel(indexers,
method=method,
tolerance=tolerance,
drop=drop)
def sel(self, indexers: dict = None, **indexers_kwargs):
# TODO: Figure out how to handle one indexer instead of start, end
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel")
for k, v in indexers.items():
if k not in self.dimensions:
raise ValueError(
f"'{k}' is not a dimension in this data stream.")
if k == "time":
if isinstance(v, slice):
start_dt = parser.parse(v.start)
end_dt = parser.parse(v.stop)
else:
start_dt, end_dt = (parser.parse(value) for value in v)
arr = getattr(self, k)
time_units = self.__default_time_units
calendar = self.__default_calendar
if 'units' in arr.attrs:
time_units = arr.attrs['units']
if 'calendar' in arr.attrs:
calendar = arr.attrs['calendar']
self.__time_filter, _, _ = xr.coding.times.encode_cf_datetime(
[start_dt, end_dt], time_units, calendar)
# Peforms initial time range check from time array of data
in_time_range = self._time_range_check(arr, start_dt, end_dt)
if not in_time_range:
return self
indexers[k] = self.__time_filter
pos_indexes = self._get_dim_indexers(indexers)
data_vars = self._create_dataset_dict(pos_indexes)
self._create_dataset(data_vars)
return self
def sel(
self, indexers: Mapping[Hashable, Any] = None, **indexers_kwargs: Any
) -> Union[xr.Dataset, xr.DataArray]:
"""Selection based on a ball tree index.
The index must have been already built using `xoak.set_index()`.
It behaves mostly like :meth:`xarray.Dataset.sel` and
:meth:`xarray.DataArray.sel` methods, with some limitations:
- Orthogonal indexing is not supported
- For vectorized (point-wise) indexing, you need to supply xarray
objects
- Use it for nearest neighbor lookup only (it implicitly
assumes method="nearest")
This triggers :func:`dask.compute` if the given indexers and/or the index
coordinates are chunked.
"""
if not getattr(self, '_index', False):
raise ValueError(
'The index(es) has/have not been built yet. Call `.xoak.set_index()` first'
)
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'xoak.sel')
indices = self._query(indexers)
if not isinstance(indices, np.ndarray):
# TODO: remove (see todo below)
indices = indices.compute()
pos_indexers = self._get_pos_indexers(indices, indexers)
# TODO: issue in xarray. 1-dimensional xarray.Variables are always considered
# as OuterIndexer, while we want here VectorizedIndexer
# This would also allow lazy selection
result = self._xarray_obj.isel(indexers=pos_indexers)
return result
def sel(self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs):
"""Wrap Dataset.sel to handle units."""
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'sel')
indexers = _reassign_quantity_indexer(self._dataset, indexers)
return self._dataset.sel(indexers, method=method, tolerance=tolerance, drop=drop)
