def index_keep_dimensions(data: xr.DataArray,
indexers: Mapping[Hashable, Union[int, slice]],
by_pos: bool=False
) -> xr.DataArray:
"""Takes an xarray and key to index it. Indexes then adds back in lost dimensions"""
# store original dims
original_dims = data.dims
# index
if by_pos:
data = data.isel(indexers)
else:
data = data.sel(indexers)
# find missing dims
missing_dims = set(original_dims) - set(data.dims)
# Add back in missing dims
data = data.expand_dims(tuple(missing_dims))
# When the selection removes a dimension, xarray.expand_dims does not expand the non-indexed
# dimensions that were removed. For example, if one selects only a single zplane, it reduces
# the z physical coordinate to a coordinate scalar, and not an array of size 1. This hack
# restores the dependent axes to arrays so they can be indexed.
for primary_axis, dependent_axis in (
(Axes.X, Coordinates.X),
(Axes.Y, Coordinates.Y),
(Axes.ZPLANE, Coordinates.Z),
):
if primary_axis.value in missing_dims and is_scalar(data[dependent_axis.value]):
data[dependent_axis.value] = xr.DataArray(
np.array([data[dependent_axis.value]]),
dims=primary_axis.value)
# Reorder to correct format
return data.transpose(*original_dims)
def assert_duckarray_equal(x, y, err_msg="", verbose=True):
"""Like `np.testing.assert_array_equal`, but for duckarrays"""
__tracebackhide__ = True
if not utils.is_duck_array(x) and not utils.is_scalar(x):
x = np.asarray(x)
if not utils.is_duck_array(y) and not utils.is_scalar(y):
y = np.asarray(y)
if (utils.is_duck_array(x)
and utils.is_scalar(y)) or (utils.is_scalar(x)
and utils.is_duck_array(y)):
equiv = (x == y).all()
else:
equiv = duck_array_ops.array_equiv(x, y)
assert equiv, _format_message(x, y, err_msg=err_msg, verbose=verbose)
def __contains__(self, key):
"""Adapted from
pandas.tseries.base.DatetimeIndexOpsMixin.__contains__"""
try:
result = self.get_loc(key)
return (is_scalar(result) or type(result) == slice or
(isinstance(result, np.ndarray) and result.size))
except (KeyError, TypeError, ValueError):
return False
def _get_indexer(self, key):
"""Get indexer for rasterio array.
Parameter
---------
key: tuple of int
Returns
-------
band_key: an indexer for the 1st dimension
window: two tuples. Each consists of (start, stop).
squeeze_axis: axes to be squeezed
np_ind: indexer for loaded numpy array
See also
--------
indexing.decompose_indexer
"""
if len(key) != 3:
raise RioXarrayError("rasterio datasets should always be 3D")
# bands cannot be windowed but they can be listed
band_key = key[0]
np_inds = []
# bands (axis=0) cannot be windowed but they can be listed
if isinstance(band_key, slice):
start, stop, step = band_key.indices(self.shape[0])
band_key = np.arange(start, stop, step)
# be sure we give out a list
band_key = (np.asarray(band_key) + 1).tolist()
if isinstance(band_key, list): # if band_key is not a scalar
np_inds.append(slice(None))
# but other dims can only be windowed
window = []
squeeze_axis = []
for i, (k, n) in enumerate(zip(key[1:], self.shape[1:])):
if isinstance(k, slice):
# step is always positive. see indexing.decompose_indexer
start, stop, step = k.indices(n)
np_inds.append(slice(None, None, step))
elif is_scalar(k):
# windowed operations will always return an array
# we will have to squeeze it later
squeeze_axis.append(-(2 - i))
start = k
stop = k + 1
else:
start, stop = np.min(k), np.max(k) + 1
np_inds.append(k - start)
window.append((start, stop))
if isinstance(key[1], np.ndarray) and isinstance(key[2], np.ndarray):
# do outer-style indexing
np_inds[-2:] = np.ix_(*np_inds[-2:])
return band_key, tuple(window), tuple(squeeze_axis), tuple(np_inds)
def __contains__(self, key):
"""Adapted from
pandas.tseries.base.DatetimeIndexOpsMixin.__contains__"""
try:
result = self.get_loc(key)
return (is_scalar(result) or type(result) == slice
or (isinstance(result, np.ndarray) and result.size))
except (KeyError, TypeError, ValueError):
return False
def drop(self, labels, dim=None, inplace=False):
"""Drop variables or index labels from this dataset. Based on xarray.dataset.drop, but adds inplace option.
Parameters
----------
labels : scalar or list of scalars
Name(s) of variables or index labels to drop.
dim : None or str, optional
Dimension along which to drop index labels. By default (if
``dim is None``), drops variables rather than index labels.
inplace : whether the original dataset should be modified or a new one created
Returns
-------
dropped : Dataset (self if inplace=True)
"""
if utils.is_scalar(labels):
labels = [labels]
if dim is None:
self._assert_all_in_dataset(labels)
drop = set(labels)
variables = OrderedDict(
(k, v) for k, v in iteritems(self._variables) if k not in drop)
coord_names = set(k for k in self._coord_names if k in variables)
result = self._replace_vars_and_dims(variables,
coord_names,
inplace=inplace)
else:
try:
index = self.indexes[dim]
except KeyError:
raise ValueError('dimension %r does not have coordinate labels' %
dim)
new_index = index.drop(labels)
result = self.loc[{dim: new_index}]
return self if inplace else result
def test_dask_array_is_scalar():
# regression test for GH1684
import dask.array as da
y = da.arange(8, chunks=4)
assert not utils.is_scalar(y)
def test_dask_array_is_scalar():
# regression test for GH1684
import dask.array as da
y = da.arange(8, chunks=4)
assert not utils.is_scalar(y)
def _atleast_1d(inp):
if utils.is_scalar(inp):
inp = (inp, )
return inp
