Merge pull request #2 from shoyer/indexing_broadcasting

Minor cleanup for Dataset.__getitem__

Author: fujiisoup

xarray/core/computation.py

@@ -136,6 +136,23 @@ def build_output_coords(
         signature,                 # type: _UFuncSignature
         exclude_dims=frozenset(),  # type: set
 ):
+    """Build output coordinates for an operation.
+
+    Parameters
+    ----------
+    args : list
+        List of raw operation arguments. Any valid types for xarray operations
+        are OK, e.g., scalars, Variable, DataArray, Dataset.
+    signature : _UfuncSignature
+        Core dimensions signature for the operation.
+    exclude_dims : optional set
+        Dimensions excluded from the operation. Coordinates along these
+        dimensions are dropped.
+
+    Returns
+    -------
+    OrderedDict of Variable objects with merged coordinates.
+    """
     # type: (...) -> List[OrderedDict[Any, Variable]]
     input_coords = _get_coord_variables(args)
 

xarray/core/dataset.py

@@ -1135,23 +1135,29 @@ def isel(self, drop=False, **indexers):
             raise ValueError("dimensions %r do not exist" % invalid)
 
         # extract new coordinates from indexers
-        variables = merge_variables([v._coords for _, v in
-                                     iteritems(indexers)
-                                     if isinstance(v, DataArray)],
-                                    compat='identical')
-
+        # we don't need to call align() explicitly, because merge_variables
+        # already checks for exact alignment between dimension coordinates
+        variables = merge_variables([v._coords for v in indexers.values()
+                                     if isinstance(v, DataArray)])
         coord_names = set(self._coord_names) | set(variables)
 
-        # a tuple, e.g. (('x', ), [0, 1]), is converted to Variable
         # all indexers should be int, slice, np.ndarrays, or Variable
-        indexers = [(k, Variable(dims=v[0], data=v[1]) if isinstance(v, tuple)
-                     else v if isinstance(v, integer_types + (slice, Variable))
-                     else v.variable if isinstance(v, DataArray)
-                     else np.asarray(v))
-                    for k, v in iteritems(indexers)]
+        indexers_list = []
+        for k, v in iteritems(indexers):
+            if isinstance(v, integer_types + (slice, Variable)):
+                pass
+            elif isinstance(v, DataArray):
+                v = v.variable
+            elif isinstance(v, tuple):
+                v = as_variable(v)
+            elif isinstance(v, Dataset):
+                raise TypeError('cannot use a Dataset as an indexer')
+            else:
+                v = np.asarray(v)
+            indexers_list.append((k, v))
 
         for name, var in iteritems(self._variables):
-            var_indexers = dict((k, v) for k, v in indexers if k in var.dims)
+            var_indexers = {k: v for k, v in indexers_list if k in var.dims}
             new_var = var.isel(**var_indexers)
             if not (drop and name in var_indexers):
                 variables[name] = new_var

xarray/core/indexing.py

@@ -7,6 +7,7 @@
 import pandas as pd
 
 from . import utils
+from .npcompat import moveaxis
 from .pycompat import (iteritems, range, integer_types, dask_array_type,
                        suppress)
 from .utils import is_dict_like
@@ -314,6 +315,10 @@ class VectorizedIndexer(IndexerTuple):
     """ Tuple for vectorized indexing """
 
 
+class PointwiseIndexer(IndexerTuple):
+    """ Tuple for pointwise indexing with dask.array's vindex """
+
+
 class LazilyIndexedArray(utils.NDArrayMixin):
     """Wrap an array that handles orthogonal indexing to make indexing lazy
     """
@@ -477,10 +482,11 @@ def __init__(self, array):
         self.array = array
 
     def __getitem__(self, key):
-        if isinstance(key, VectorizedIndexer):
-            # TODO should support vindex
-            raise IndexError(
-                'dask does not support vectorized indexing : {}'.format(key))
+        # should always get PointwiseIndexer instead
+        assert not isinstance(key, VectorizedIndexer)
+
+        if isinstance(key, PointwiseIndexer):
+            return self._getitem_pointwise(key)
 
         try:
             key = to_tuple(key)
@@ -492,6 +498,30 @@ def __getitem__(self, key):
                 value = value[(slice(None),) * axis + (subkey,)]
             return value
 
+    def _getitem_pointwise(self, key):
+        pointwise_shape, pointwise_index = next(
+           (k.shape, i) for i, k in enumerate(key)
+           if not isinstance(k, slice))
+        # dask's indexing only handles 1d arrays
+        flat_key = tuple(k if isinstance(k, slice) else k.ravel()
+                         for k in key)
+
+        if len([k for k in key if not isinstance(k, slice)]) == 1:
+            # vindex requires more than one non-slice :(
+            # but we can use normal indexing instead
+            indexed = self.array[flat_key]
+            new_shape = (indexed.shape[:pointwise_index] +
+                         pointwise_shape +
+                         indexed.shape[pointwise_index + 1:])
+            return indexed.reshape(new_shape)
+        else:
+            indexed = self.array.vindex[flat_key]
+            # vindex always moves slices to the end
+            reshaped = indexed.reshape(pointwise_shape + indexed.shape[1:])
+            # reorder dimensions to match order of appearance
+            positions = np.arange(0, len(pointwise_shape))
+            return moveaxis(reshaped, positions, positions + pointwise_index)
+
     def __setitem__(self, key, value):
         raise TypeError("this variable's data is stored in a dask array, "
                         'which does not support item assignment. To '

xarray/core/npcompat.py

@@ -1,10 +1,14 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import operator
+
 import numpy as np
 
 try:
-    from numpy import broadcast_to, stack, nanprod, nancumsum, nancumprod
+    from numpy import (broadcast_to, stack, nanprod, nancumsum, nancumprod,
+                       moveaxis)
 except ImportError:  # pragma: no cover
     # Code copied from newer versions of NumPy (v1.10 to v1.12).
     # Used under the terms of NumPy's license, see licenses/NUMPY_LICENSE.
@@ -371,3 +375,130 @@ def nancumprod(a, axis=None, dtype=None, out=None):
         """
         a, mask = _replace_nan(a, 1)
         return np.cumprod(a, axis=axis, dtype=dtype, out=out)
+
+
+    def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
+        """
+        Normalizes an axis argument into a tuple of non-negative integer axes.
+
+        This handles shorthands such as ``1`` and converts them to ``(1,)``,
+        as well as performing the handling of negative indices covered by
+        `normalize_axis_index`.
+
+        By default, this forbids axes from being specified multiple times.
+
+        Used internally by multi-axis-checking logic.
+
+        .. versionadded:: 1.13.0
+
+        Parameters
+        ----------
+        axis : int, iterable of int
+            The un-normalized index or indices of the axis.
+        ndim : int
+            The number of dimensions of the array that `axis` should be normalized
+            against.
+        argname : str, optional
+            A prefix to put before the error message, typically the name of the
+            argument.
+        allow_duplicate : bool, optional
+            If False, the default, disallow an axis from being specified twice.
+
+        Returns
+        -------
+        normalized_axes : tuple of int
+            The normalized axis index, such that `0 <= normalized_axis < ndim`
+
+        Raises
+        ------
+        AxisError
+            If any axis provided is out of range
+        ValueError
+            If an axis is repeated
+
+        See also
+        --------
+        normalize_axis_index : normalizing a single scalar axis
+        """
+        try:
+            axis = [operator.index(axis)]
+        except TypeError:
+            axis = tuple(axis)
+        axis = tuple(normalize_axis_index(ax, ndim, argname) for ax in axis)
+        if not allow_duplicate and len(set(axis)) != len(axis):
+            if argname:
+                raise ValueError('repeated axis in `{}` argument'.format(argname))
+            else:
+                raise ValueError('repeated axis')
+        return axis
+
+
+    def moveaxis(a, source, destination):
+        """
+        Move axes of an array to new positions.
+
+        Other axes remain in their original order.
+
+        .. versionadded::1.11.0
+
+        Parameters
+        ----------
+        a : np.ndarray
+            The array whose axes should be reordered.
+        source : int or sequence of int
+            Original positions of the axes to move. These must be unique.
+        destination : int or sequence of int
+            Destination positions for each of the original axes. These must also be
+            unique.
+
+        Returns
+        -------
+        result : np.ndarray
+            Array with moved axes. This array is a view of the input array.
+
+        See Also
+        --------
+        transpose: Permute the dimensions of an array.
+        swapaxes: Interchange two axes of an array.
+
+        Examples
+        --------
+
+        >>> x = np.zeros((3, 4, 5))
+        >>> np.moveaxis(x, 0, -1).shape
+        (4, 5, 3)
+        >>> np.moveaxis(x, -1, 0).shape
+        (5, 3, 4)
+
+        These all achieve the same result:
+
+        >>> np.transpose(x).shape
+        (5, 4, 3)
+        >>> np.swapaxes(x, 0, -1).shape
+        (5, 4, 3)
+        >>> np.moveaxis(x, [0, 1], [-1, -2]).shape
+        (5, 4, 3)
+        >>> np.moveaxis(x, [0, 1, 2], [-1, -2, -3]).shape
+        (5, 4, 3)
+
+        """
+        try:
+            # allow duck-array types if they define transpose
+            transpose = a.transpose
+        except AttributeError:
+            a = np.asarray(a)
+            transpose = a.transpose
+
+        source = normalize_axis_tuple(source, a.ndim, 'source')
+        destination = normalize_axis_tuple(destination, a.ndim, 'destination')
+        if len(source) != len(destination):
+            raise ValueError('`source` and `destination` arguments must have '
+                             'the same number of elements')
+
+        order = [n for n in range(a.ndim) if n not in source]
+
+        for dest, src in sorted(zip(destination, source)):
+            order.insert(dest, src)
+
+        result = transpose(order)
+        return result

xarray/core/variable.py

@@ -19,7 +19,7 @@
 from .pycompat import (basestring, OrderedDict, zip, integer_types,
                        dask_array_type)
 from .indexing import (PandasIndexAdapter, xarray_indexable, BasicIndexer,
-                       OuterIndexer, VectorizedIndexer)
+                       OuterIndexer, PointwiseIndexer, VectorizedIndexer)
 
 import xarray as xr  # only for Dataset and DataArray
 
@@ -464,8 +464,11 @@ def _nonzero(self):
                      in zip(nonzeros, self.dims))
 
     def _broadcast_indexes_advanced(self, key):
-        variables = []
+        if isinstance(self._data, dask_array_type):
+            # dask only supports a very restricted form of advanced indexing
+            return self._broadcast_indexes_dask_pointwise(key)
 
+        variables = []
         for dim, value in zip(self.dims, key):
             if isinstance(value, slice):
                 value = np.arange(*value.indices(self.sizes[dim]))
@@ -492,6 +495,46 @@ def _broadcast_indexes_advanced(self, key):
         key = VectorizedIndexer(variable.data for variable in variables)
         return dims, key
 
+    def _broadcast_indexes_dask_pointwise(self, key):
+        if any(not isinstance(k, (Variable, slice)) for k in key):
+            raise IndexError(
+                'Vectorized indexing with dask requires that all indexers are '
+                'labeled arrays or full slice objects: {}'.format(key))
+
+        if any(isinstance(k, Variable) and k.dtype.kind == 'b' for k in key):
+            raise IndexError(
+                'Vectorized indexing with dask does not support booleans: {}'
+                .format(key))
+
+        dims_set = {k.dims for k in key if isinstance(k, Variable)}
+        if len(dims_set) != 1:
+            raise IndexError(
+                'Vectorized indexing with dask requires that all labeled '
+                'arrays in the indexer have the same dimension names, but '
+                'arrays have different dimensions: {}'.format(key))
+        (unique_dims,) = dims_set
+
+        shapes_set = {k.shape for k in key if isinstance(k, Variable)}
+        if len(shapes_set) != 1:
+            # matches message in _broadcast_indexes_advanced
+            raise IndexError("Dimensions of indexers mismatch: {}".format(key))
+
+        dims = []
+        found_first_array = False
+        for k, d in zip(key, self.dims):
+            if isinstance(k, slice):
+                if d in unique_dims:
+                    raise IndexError(
+                        'Labeled arrays used in vectorized indexing with dask '
+                        'cannot reuse a sliced dimension: {}'.format(d))
+                dims.append(d)
+            elif not found_first_array:
+                dims.extend(k.dims)
+                found_first_array = True
+
+        key = PointwiseIndexer(getattr(k, 'data', k) for k in key)
+        return tuple(dims), key
+
     def __getitem__(self, key):
         """Return a new Array object whose contents are consistent with
         getting the provided key from the underlying data.
@@ -507,12 +550,12 @@ def __getitem__(self, key):
         array `x.values` directly.
         """
         dims, index_tuple = self._broadcast_indexes(key)
-        values = self._indexable_data[index_tuple]
-        if hasattr(values, 'ndim'):
-            assert values.ndim == len(dims), (values.ndim, len(dims))
+        data = self._indexable_data[index_tuple]
+        if hasattr(data, 'ndim'):
+            assert data.ndim == len(dims), (data.ndim, len(dims))
         else:
             assert len(dims) == 0, len(dims)
-        return type(self)(dims, values, self._attrs, self._encoding,
+        return type(self)(dims, data, self._attrs, self._encoding,
                           fastpath=True)
 
     def __setitem__(self, key, value):

xarray/tests/test_dataset.py

@@ -881,6 +881,8 @@ def test_isel_fancy(self):
         with self.assertRaisesRegexp(IndexError,
                                      'Dimensions of indexers mismatch'):
             data.isel(dim1=(('points'), [1, 2]), dim2=(('points'), [1, 2, 3]))
+        with self.assertRaisesRegexp(TypeError, 'cannot use a Dataset'):
+            data.isel(dim1=Dataset({'points': [1, 2]}))
 
         # test to be sure we keep around variables that were not indexed
         ds = Dataset({'x': [1, 2, 3, 4], 'y': 0})