Multi-index levels as coordinates (#947)

* make multi-index levels visible as coordinates

* make levels also visible for Dataset

* fix unnamed levels

* allow providing multi-index levels in .sel

* refactored _get_valid_indexers to get_dim_indexers

* fix broken tests

* refactored accessibility and repr of index levels

* do not allow providing both level and dim indexers in .sel

* cosmetic changes

* change signature of Coordinate.__init__

* check for uniqueness of multi-index level names

* no need to check for uniqueness of level names in _level_coords

* rewritten checking uniqueness of multi-index level names

* fix adding coords/vars with the same name than a multi-index level

* check for level/var name conflicts in one place

* cosmetic changes

* fix Coordinate -> IndexVariable

* fix col width when formatting multi-index levels

* add tests for IndexVariable new methods and indexing

* fix bug in assert_unique_multiindex_level_names

* add tests for Dataset

* fix appveyor tests

* add tests for DataArray

* add docs

* review changes

* remove name argument of IndexVariable

Author: Benoit Bovy

doc/data-structures.rst

@@ -115,10 +115,6 @@ If you create a ``DataArray`` by supplying a pandas
     df
     xr.DataArray(df)
 
-Xarray supports labeling coordinate values with a :py:class:`pandas.MultiIndex`.
-While it handles multi-indexes with unnamed levels, it is recommended that you
-explicitly set the names of the levels.
-
 DataArray properties
 ~~~~~~~~~~~~~~~~~~~~
 
@@ -532,6 +528,41 @@ dimension and whose the values are ``Index`` objects:
 
     ds.indexes
 
+MultiIndex coordinates
+~~~~~~~~~~~~~~~~~~~~~~
+
+Xarray supports labeling coordinate values with a :py:class:`pandas.MultiIndex`:
+
+.. ipython:: python
+
+    midx = pd.MultiIndex.from_arrays([['R', 'R', 'V', 'V'], [.1, .2, .7, .9]],
+                                     names=('band', 'wn'))
+    mda = xr.DataArray(np.random.rand(4), coords={'spec': midx}, dims='spec')
+    mda
+
+For convenience multi-index levels are directly accessible as "virtual" or
+"derived" coordinates (marked by ``-`` when printing a dataset or data array):
+
+.. ipython:: python
+
+     mda['band']
+     mda.wn
+
+Indexing with multi-index levels is also possible using the ``sel`` method
+(see :ref:`multi-level indexing`).
+
+Unlike other coordinates, "virtual" level coordinates are not stored in
+the ``coords`` attribute of ``DataArray`` and ``Dataset`` objects
+(although they are shown when printing the ``coords`` attribute).
+Consequently, most of the coordinates related methods don't apply for them.
+It also can't be used to replace one particular level.
+
+Because in a ``DataArray`` or ``Dataset`` object each multi-index level is
+accessible as a "virtual" coordinate, its name must not conflict with the names
+of the other levels, coordinates and data variables of the same object.
+Even though Xarray set default names for multi-indexes with unnamed levels,
+it is recommended that you explicitly set the names of the levels.
+
 .. [1] Latitude and longitude are 2D arrays because the dataset uses
    `projected coordinates`__. ``reference_time`` refers to the reference time
    at which the forecast was made, rather than ``time`` which is the valid time

doc/indexing.rst

@@ -325,11 +325,25 @@ Additionally, xarray supports dictionaries:
 .. ipython:: python
 
    mda.sel(x={'one': 'a', 'two': 0})
-   mda.loc[{'one': 'a'}, ...]
+
+For convenience, ``sel`` also accepts multi-index levels directly
+as keyword arguments:
+
+.. ipython:: python
+
+   mda.sel(one='a', two=0)
+
+Note that using ``sel`` it is not possible to mix a dimension
+indexer with level indexers for that dimension
+(e.g., ``mda.sel(x={'one': 'a'}, two=0)`` will raise a ``ValueError``).
 
 Like pandas, xarray handles partial selection on multi-index (level drop).
-As shown in the last example above, it also renames the dimension / coordinate
-when the multi-index is reduced to a single index.
+As shown below, it also renames the dimension / coordinate when the
+multi-index is reduced to a single index.
+
+.. ipython:: python
+
+   mda.loc[{'one': 'a'}, ...]
 
 Unlike pandas, xarray does not guess whether you provide index levels or
 dimensions when using ``loc`` in some ambiguous cases. For example, for

doc/whats-new.rst

@@ -50,6 +50,13 @@ By `Robin Wilson <https://github.com/robintw>`_.
   and deals with names to ensure a perfect 'roundtrip' capability.
   By `Robin Wilson <https://github.com/robintw`_.
 
+- Multi-index levels are now accessible as "virtual" coordinate variables,
+  e.g., ``ds['time']`` can pull out the ``'time'`` level of a multi-index
+  (see :ref:`coordinates`). ``sel`` also accepts providing multi-index levels
+  as keyword arguments, e.g., ``ds.sel(time='2000-01')``
+  (see :ref:`multi-level indexing`).
+  By `Benoit Bovy <https://github.com/benbovy>`_.
+
 Bug fixes
 ~~~~~~~~~
 

xarray/core/coordinates.py

@@ -215,6 +215,27 @@ def __delitem__(self, key):
         del self._data._coords[key]
 
 
+class DataArrayLevelCoordinates(AbstractCoordinates):
+    """Dictionary like container for DataArray MultiIndex level coordinates.
+
+    Used for attribute style lookup. Not returned directly by any
+    public methods.
+    """
+    def __init__(self, dataarray):
+        self._data = dataarray
+
+    @property
+    def _names(self):
+        return set(self._data._level_coords)
+
+    @property
+    def variables(self):
+        level_coords = OrderedDict(
+            (k, self._data[v].variable.get_level_variable(k))
+             for k, v in self._data._level_coords.items())
+        return Frozen(level_coords)
+
+
 class Indexes(Mapping, formatting.ReprMixin):
     """Ordered Mapping[str, pandas.Index] for xarray objects.
     """

xarray/core/dataarray.py

@@ -14,11 +14,13 @@
 from . import utils
 from .alignment import align
 from .common import AbstractArray, BaseDataObject, squeeze
-from .coordinates import DataArrayCoordinates, Indexes
+from .coordinates import (DataArrayCoordinates, DataArrayLevelCoordinates,
+                          Indexes)
 from .dataset import Dataset
 from .pycompat import iteritems, basestring, OrderedDict, zip
 from .variable import (as_variable, Variable, as_compatible_data, IndexVariable,
-                       default_index_coordinate)
+                       default_index_coordinate,
+                       assert_unique_multiindex_level_names)
 from .formatting import format_item
 
 
@@ -82,6 +84,8 @@ def _infer_coords_and_dims(shape, coords, dims):
                                  'length %s on the data but length %s on '
                                  'coordinate %r' % (d, sizes[d], s, k))
 
+    assert_unique_multiindex_level_names(new_coords)
+
     return new_coords, dims
 
 
@@ -421,14 +425,29 @@ def _item_key_to_dict(self, key):
             key = indexing.expanded_indexer(key, self.ndim)
             return dict(zip(self.dims, key))
 
+    @property
+    def _level_coords(self):
+        """Return a mapping of all MultiIndex levels and their corresponding
+        coordinate name.
+        """
+        level_coords = OrderedDict()
+        for cname, var in self._coords.items():
+            if var.ndim == 1:
+                level_names = var.to_index_variable().level_names
+                if level_names is not None:
+                    dim, = var.dims
+                    level_coords.update({lname: dim for lname in level_names})
+        return level_coords
+
     def __getitem__(self, key):
         if isinstance(key, basestring):
             from .dataset import _get_virtual_variable
 
             try:
                 var = self._coords[key]
             except KeyError:
-                _, key, var = _get_virtual_variable(self._coords, key)
+                _, key, var = _get_virtual_variable(
+                    self._coords, key, self._level_coords)
 
             return self._replace_maybe_drop_dims(var, name=key)
         else:
@@ -448,7 +467,7 @@ def __delitem__(self, key):
     @property
     def _attr_sources(self):
         """List of places to look-up items for attribute-style access"""
-        return [self.coords, self.attrs]
+        return [self.coords, DataArrayLevelCoordinates(self), self.attrs]
 
     def __contains__(self, key):
         return key in self._coords

xarray/core/dataset.py

@@ -33,34 +33,48 @@
                              'quarter']
 
 
-def _get_virtual_variable(variables, key):
-    """Get a virtual variable (e.g., 'time.year') from a dict of
-    xarray.Variable objects (if possible)
+def _get_virtual_variable(variables, key, level_vars={}):
+    """Get a virtual variable (e.g., 'time.year' or a MultiIndex level)
+    from a dict of xarray.Variable objects (if possible)
     """
     if not isinstance(key, basestring):
         raise KeyError(key)
 
     split_key = key.split('.', 1)
-    if len(split_key) != 2:
+    if len(split_key) == 2:
+        ref_name, var_name = split_key
+    elif len(split_key) == 1:
+        ref_name, var_name = key, None
+    else:
         raise KeyError(key)
 
-    ref_name, var_name = split_key
-    ref_var = variables[ref_name]
-    if ref_var.ndim == 1:
-        date = ref_var.to_index()
-    elif ref_var.ndim == 0:
-        date = pd.Timestamp(ref_var.values)
+    if ref_name in level_vars:
+        dim_var = variables[level_vars[ref_name]]
+        ref_var = dim_var.to_index_variable().get_level_variable(ref_name)
     else:
-        raise KeyError(key)
+        ref_var = variables[ref_name]
 
-    if var_name == 'season':
-        # TODO: move 'season' into pandas itself
-        seasons = np.array(['DJF', 'MAM', 'JJA', 'SON'])
-        month = date.month
-        data = seasons[(month // 3) % 4]
+    if var_name is None:
+        virtual_var = ref_var
+        var_name = key
     else:
-        data = getattr(date, var_name)
-    return ref_name, var_name, Variable(ref_var.dims, data)
+        if ref_var.ndim == 1:
+            date = ref_var.to_index()
+        elif ref_var.ndim == 0:
+            date = pd.Timestamp(ref_var.values)
+        else:
+            raise KeyError(key)
+
+        if var_name == 'season':
+            # TODO: move 'season' into pandas itself
+            seasons = np.array(['DJF', 'MAM', 'JJA', 'SON'])
+            month = date.month
+            data = seasons[(month // 3) % 4]
+        else:
+            data = getattr(date, var_name)
+        virtual_var = Variable(ref_var.dims, data)
+
+    return ref_name, var_name, virtual_var
 
 
 def calculate_dimensions(variables):
@@ -411,6 +425,21 @@ def _subset_with_all_valid_coords(self, variables, coord_names, attrs):
 
         return self._construct_direct(variables, coord_names, dims, attrs)
 
+    @property
+    def _level_coords(self):
+        """Return a mapping of all MultiIndex levels and their corresponding
+        coordinate name.
+        """
+        level_coords = OrderedDict()
+        for cname in self._coord_names:
+            var = self.variables[cname]
+            if var.ndim == 1:
+                level_names = var.to_index_variable().level_names
+                if level_names is not None:
+                    dim, = var.dims
+                    level_coords.update({lname: dim for lname in level_names})
+        return level_coords
+
     def _copy_listed(self, names):
         """Create a new Dataset with the listed variables from this dataset and
         the all relevant coordinates. Skips all validation.
@@ -423,7 +452,7 @@ def _copy_listed(self, names):
                 variables[name] = self._variables[name]
             except KeyError:
                 ref_name, var_name, var = _get_virtual_variable(
-                    self._variables, name)
+                    self._variables, name, self._level_coords)
                 variables[var_name] = var
                 if ref_name in self._coord_names:
                     coord_names.add(var_name)
@@ -439,7 +468,8 @@ def _construct_dataarray(self, name):
         try:
             variable = self._variables[name]
         except KeyError:
-            _, name, variable = _get_virtual_variable(self._variables, name)
+            _, name, variable = _get_virtual_variable(
+                self._variables, name, self._level_coords)
 
         coords = OrderedDict()
         needed_dims = set(variable.dims)
@@ -508,6 +538,7 @@ def __setitem__(self, key, value):
         if utils.is_dict_like(key):
             raise NotImplementedError('cannot yet use a dictionary as a key '
                                       'to set Dataset values')
+
         self.update({key: value})
 
     def __delitem__(self, key):