Merge pull request #291 from shoyer/faster-isel

Support using dictionaries for labeled indexing

Author: shoyer

doc/api.rst

@@ -16,7 +16,6 @@ Top-level functions
 
    align
    concat
-   decode_cf
 
 Dataset
 =======
@@ -29,6 +28,7 @@ Creating a dataset
 
    Dataset
    open_dataset
+   decode_cf
 
 Attributes
 ----------
@@ -85,6 +85,7 @@ Indexing
 .. autosummary::
    :toctree: generated/
 
+   Dataset.loc
    Dataset.isel
    Dataset.sel
    Dataset.squeeze

doc/indexing.rst

@@ -28,8 +28,8 @@ By name          By integer   ``arr.isel(space=0)``   ``ds.isel(space=0)``
 By name          By label     ``arr.sel(space='IA')`` ``ds.sel(space='IA')``
 ================ ============ ======================= ======================
 
-Array indexing
---------------
+Positional indexing
+-------------------
 
 Indexing a :py:class:`~xray.DataArray` directly works (mostly) just like it
 does for numpy arrays, except that the returned object is always another
@@ -70,16 +70,29 @@ Indexing with labeled dimensions
 --------------------------------
 
 With labeled dimensions, we do not have to rely on dimension order and can
-use them explicitly to slice data with the :py:meth:`~xray.DataArray.sel`
-and :py:meth:`~xray.DataArray.isel` methods:
+use them explicitly to slice data. There are two ways to do this:
 
-.. ipython:: python
+1. Use a dictionary as the argument for array positional or label based array
+   indexing:
+
+    .. ipython:: python
+
+        # index by integer array indices
+        arr[dict(space=0, time=slice(None, 2))]
+
+        # index by dimension coordinate labels
+        arr.loc[dict(time=slice('2000-01-01', '2000-01-02'))]
+
+2. Use the :py:meth:`~xray.DataArray.sel` and :py:meth:`~xray.DataArray.isel`
+   convenience methods:
 
-    # index by integer array indices
-    arr.isel(space=0, time=slice(None, 2))
+    .. ipython:: python
 
-    # index by dimension coordinate labels
-    arr.sel(time=slice('2000-01-01', '2000-01-02'))
+        # index by integer array indices
+        arr.isel(space=0, time=slice(None, 2))
+
+        # index by dimension coordinate labels
+        arr.sel(time=slice('2000-01-01', '2000-01-02'))
 
 The arguments to these methods can be any objects that could index the array
 along the dimension given by the keyword, e.g., labels for an individual value,
@@ -88,10 +101,8 @@ Python :py:func:`slice` objects or 1-dimensional arrays.
 .. note::
 
     We would love to be able to do indexing with labeled dimension names inside
-    brackets, but Python `does yet not support`__ indexing with keyword
-    arguments like ``arr[space=0]``. One alternative we are considering is
-    allowing for indexing with a dictionary, ``arr[{'space': 0}]``
-    (see :issue:`187`.
+    brackets, but unfortunately, Python `does yet not support`__ indexing with
+    keyword arguments like ``arr[space=0]``
 
 __ http://legacy.python.org/dev/peps/pep-0472/
 
@@ -100,17 +111,14 @@ __ http://legacy.python.org/dev/peps/pep-0472/
     Do not try to assign values when using ``isel`` or ``sel``::
 
         # DO NOT do this
-        arr.isel(space='0') = 0
+        arr.isel(space=0) = 0
 
     Depending on whether the underlying numpy indexing returns a copy or a
     view, the method will fail, and when it fails, **it will fail
-    silently**. Until we support indexing with dictionaries (see the note
-    above), you should explicitly construct a tuple to do positional indexing
-    if you want to do assignment with labeled dimensions::
+    silently**. Instead, you should use normal index assignment::
 
-        # this is safer
-        indexer = tuple(0 if d == 'space' else slice(None) for d in arr.dims)
-        arr[indexer] = 0
+        # this is safe
+        arr[dict(space=0)] = 0
 
 Dataset indexing
 ----------------
@@ -126,7 +134,15 @@ simultaneously, returning a new dataset:
 
 Positional indexing on a dataset is not supported because the ordering of
 dimensions in a dataset is somewhat ambiguous (it can vary between different
-arrays).
+arrays). However, you can do normal indexing with labeled dimensions:
+
+.. ipython:: python
+
+    ds[dict(space=[0], time=[0])]
+    ds.loc[dict(time='2000-01-01')]
+
+Using indexing to *assign* values to a subset of dataset (e.g.,
+``ds[dict(space=0)] = 1``) is not yet supported.
 
 Indexing details
 ----------------

xray/__init__.py

@@ -3,6 +3,6 @@
 from .core.dataset import Dataset, open_dataset
 from .core.dataarray import DataArray
 
-from .conventions import cf_decode
+from .conventions import decode_cf
 
 from .version import version as __version__

xray/conventions.py

@@ -638,9 +638,10 @@ def stackable(dim):
     return new_vars, attributes, coord_names
 
 
-def cf_decode(obj, concat_characters=True, mask_and_scale=True,
+def decode_cf(obj, concat_characters=True, mask_and_scale=True,
               decode_times=True, decode_coords=True):
-    """Decode the given object according to CF conventions into a new Dataset.
+    """Decode the given Dataset or Datastore according to CF conventions into
+    a new Dataset.
 
     Parameters
     ----------

xray/core/dataarray.py

@@ -62,12 +62,20 @@ def __init__(self, data_array):
         self.data_array = data_array
 
     def _remap_key(self, key):
-        label_indexers = self.data_array._key_to_indexers(key)
-        indexers = []
-        for dim, label in iteritems(label_indexers):
+        def lookup_positions(dim, labels):
             index = self.data_array.indexes[dim]
-            indexers.append(indexing.convert_label_indexer(index, label))
-        return tuple(indexers)
+            return indexing.convert_label_indexer(index, labels)
+
+        if utils.is_dict_like(key):
+            return dict((dim, lookup_positions(dim, labels))
+                        for dim, labels in iteritems(key))
+        else:
+            if not isinstance(key, tuple):
+                key = (key,)
+            # note: it's OK if there are fewer keys than dimensions: zip will
+            # finish early in that case (we don't need to insert colons)
+            return tuple(lookup_positions(dim, labels) for dim, labels
+                         in zip(self.data_array.dims, key))
 
     def __getitem__(self, key):
         return self.data_array[self._remap_key(key)]
@@ -326,16 +334,19 @@ def dimensions(self):
         utils.alias_warning('dimensions', 'dims')
         return self.dims
 
-    def _key_to_indexers(self, key):
-        return OrderedDict(
-            zip(self.dims, indexing.expanded_indexer(key, self.ndim)))
+    def _item_key_to_dict(self, key):
+        if utils.is_dict_like(key):
+            return key
+        else:
+            key = indexing.expanded_indexer(key, self.ndim)
+            return dict(zip(self.dims, key))
 
     def __getitem__(self, key):
         if isinstance(key, basestring):
             return self.coords[key]
         else:
             # orthogonal array indexing
-            return self.isel(**self._key_to_indexers(key))
+            return self.isel(**self._item_key_to_dict(key))
 
     def __setitem__(self, key, value):
         if isinstance(key, basestring):
@@ -352,7 +363,7 @@ def __contains__(self, key):
 
     @property
     def loc(self):
-        """Attribute for location based indexing like pandas..
+        """Attribute for location based indexing like pandas.
         """
         return _LocIndexer(self)
 

xray/core/dataset.py

@@ -85,7 +85,7 @@ def open_dataset(filename_or_obj, decode_cf=True, mask_and_scale=True,
         store = backends.ScipyDataStore(filename_or_obj)
 
     if decode_cf:
-        return conventions.cf_decode(
+        return conventions.decode_cf(
             store, mask_and_scale=mask_and_scale,
             decode_times=decode_times, concat_characters=concat_characters,
             decode_coords=decode_coords)
@@ -292,6 +292,16 @@ def __repr__(self):
         return formatting.vars_repr(self)
 
 
+class _LocIndexer(object):
+    def __init__(self, dataset):
+        self.dataset = dataset
+
+    def __getitem__(self, key):
+        if not utils.is_dict_like(key):
+            raise TypeError('can only lookup dictionaries from Dataset.loc')
+        return self.dataset.sel(**key)
+
+
 class Dataset(Mapping, common.ImplementsDatasetReduce):
     """A multi-dimensional, in memory, array database.
 
@@ -606,6 +616,13 @@ def __len__(self):
     def __iter__(self):
         return iter(self._arrays)
 
+    @property
+    def loc(self):
+        """Attribute for location based indexing. Only supports __getitem__,
+        and only when the key is a dict of the form {dim: labels}.
+        """
+        return _LocIndexer(self)
+
     @property
     def virtual_variables(self):
         """A frozenset of names that don't exist in this dataset but for which
@@ -633,6 +650,9 @@ def __getitem__(self, key):
         """
         from .dataarray import DataArray
 
+        if utils.is_dict_like(key):
+            return self.isel(**key)
+
         key = np.asarray(key)
         if key.ndim == 0:
             return DataArray._new_from_dataset(self, key.item())
@@ -650,6 +670,9 @@ def __setitem__(self, key, value):
         ``(dims, data[, attrs])``), add it to this dataset as a new
         variable.
         """
+        if utils.is_dict_like(key):
+            raise NotImplementedError('cannot yet use a dictionary as a key '
+                                      'to set Dataset values')
         self.merge({key: value}, inplace=True, overwrite_vars=[key])
 
     def __delitem__(self, key):
@@ -967,21 +990,23 @@ def reindex_like(self, other, copy=True):
         """
         return self.reindex(copy=copy, **other.indexes)
 
-    def reindex(self, copy=True, **indexers):
+    def reindex(self, indexers=None, copy=True, **kw_indexers):
         """Conform this object onto a new set of indexes, filling in
         missing values with NaN.
 
         Parameters
         ----------
-        copy : bool, optional
-            If `copy=True`, the returned dataset contains only copied
-            variables. If `copy=False` and no reindexing is required then
-            original variables from this dataset are returned.
-        **indexers : dict
+        indexers : dict. optional
             Dictionary with keys given by dimension names and values given by
             arrays of coordinates tick labels. Any mis-matched coordinate values
             will be filled in with NaN, and any mis-matched dimension names will
             simply be ignored.
+        copy : bool, optional
+            If `copy=True`, the returned dataset contains only copied
+            variables. If `copy=False` and no reindexing is required then
+            original variables from this dataset are returned.
+        **kw_indexers : optional
+            Keyword arguments in the same form as ``indexers``.
 
         Returns
         -------
@@ -993,6 +1018,8 @@ def reindex(self, copy=True, **indexers):
         Dataset.reindex_like
         align
         """
+        indexers = utils.combine_pos_and_kw_args(indexers, kw_indexers,
+                                                 'reindex')
         if not indexers:
             # shortcut
             return self.copy(deep=True) if copy else self

xray/core/utils.py

@@ -189,6 +189,19 @@ def is_dict_like(value):
     return hasattr(value, '__getitem__') and hasattr(value, 'keys')
 
 
+def combine_pos_and_kw_args(pos_kwargs, kw_kwargs, func_name):
+    if pos_kwargs is not None:
+        if not is_dict_like(pos_kwargs):
+            raise ValueError('the first argument to .%s must be a dictionary'
+                             % func_name)
+        if kw_kwargs:
+            raise ValueError('cannot specify both keyword and positional '
+                             'arguments to .%s' % func_name)
+        return pos_kwargs
+    else:
+        return kw_kwargs
+
+
 def is_scalar(value):
     """np.isscalar only work on primitive numeric types and (bizarrely)
     excludes 0-d ndarrays; this version does more comprehensive checks

xray/core/variable.py

@@ -357,6 +357,12 @@ def _parse_dimensions(self, dims):
     def dims(self, value):
         self._dims = self._parse_dimensions(value)
 
+    def _item_key_to_tuple(self, key):
+        if utils.is_dict_like(key):
+            return tuple(key.get(dim, slice(None)) for dim in self.dims)
+        else:
+            return key
+
     def __getitem__(self, key):
         """Return a new Array object whose contents are consistent with
         getting the provided key from the underlying data.
@@ -374,9 +380,10 @@ def __getitem__(self, key):
         If you really want to do indexing like `x[x > 0]`, manipulate the numpy
         array `x.values` directly.
         """
+        key = self._item_key_to_tuple(key)
         key = indexing.expanded_indexer(key, self.ndim)
         dims = [dim for k, dim in zip(key, self.dims)
-                      if not isinstance(k, (int, np.integer))]
+                if not isinstance(k, (int, np.integer))]
         values = self._data[key]
         # orthogonal indexing should ensure the dimensionality is consistent
         if hasattr(values, 'ndim'):
@@ -391,6 +398,7 @@ def __setitem__(self, key, value):
 
         See __getitem__ for more details.
         """
+        key = self._item_key_to_tuple(key)
         self._data_cached()[key] = value
 
     @property
@@ -817,6 +825,7 @@ def __init__(self, name, data, attrs=None, encoding=None):
                              type(self).__name__)
 
     def __getitem__(self, key):
+        key = self._item_key_to_tuple(key)
         values = self._data[key]
         if not hasattr(values, 'ndim') or values.ndim == 0:
             return Variable((), values, self.attrs, self.encoding)

xray/test/__init__.py

@@ -97,9 +97,8 @@ def assertDatasetIdentical(self, d1, d2):
         assert d1.identical(d2), (d1, d2)
 
     def assertDatasetAllClose(self, d1, d2, rtol=1e-05, atol=1e-08):
-        # for now, *don't* check coordinates vs variables
-        self.assertEqual(sorted(d1, key=str),
-                         sorted(d2, key=str))
+        self.assertEqual(sorted(d1, key=str), sorted(d2, key=str))
+        self.assertItemsEqual(d1.coords, d2.coords)
         for k in d1:
             v1 = d1._arrays[k]
             v2 = d2._arrays[k]