Fall back sel -> isel if no index, add/fix some tests

Author: shoyer

doc/whats-new.rst

@@ -95,6 +95,10 @@ Bug fixes
 - ``.where()`` and ``.fillna()`` now preserve attributes(:issue:`1009`).
   By `Fabien Maussion <https://github.com/fmaussion>`_.
 
+- Fixed accessing coordinate variables with non-string names from ``.coords``
+  (:issue:`TBD`).
+  By `Stephan Hoyer <https://github.com/shoyer>`_.
+
 .. _whats-new.0.8.2:
 
 v0.8.2 (18 August 2016)

xarray/core/alignment.py

@@ -73,6 +73,12 @@ def align(*objects, **kwargs):
     -------
     aligned : same as *objects
         Tuple of objects with aligned coordinates.
+
+    Raises
+    ------
+    ValueError
+        If any dimensions without labels on the arguments have different sizes,
+        or a different size than the size of the aligned dimension labels.
     """
     join = kwargs.pop('join', 'inner')
     copy = kwargs.pop('copy', True)
@@ -87,19 +93,19 @@ def align(*objects, **kwargs):
                         % list(kwargs))
 
     all_indexes = defaultdict(list)
-    unlabeled_dim_sizes = defaultdict(list)
+    unlabeled_dim_sizes = defaultdict(set)
     for obj in objects:
         for dim in obj.dims:
             if dim not in exclude:
                 try:
                     index = obj.indexes[dim]
                 except KeyError:
                     size = get_size(obj, dim)
-                    unlabeled_dim_sizes[dim].append(size)
+                    unlabeled_dim_sizes[dim].add(size)
                 else:
                     all_indexes[dim].append(index)
 
-    # We don't join over dimensions with all equal indexes for two reasons:
+    # We don't reindex over dimensions with all equal indexes for two reasons:
     # - It's faster for the usual case (already aligned objects).
     # - It ensures it's possible to do operations that don't require alignment
     #   on indexes with duplicate values (which cannot be reindexed with
@@ -114,20 +120,34 @@ def align(*objects, **kwargs):
         else:
             if any(not matching_indexes[0].equals(other)
                    for other in matching_indexes[1:]):
-                joined_indexes[name] = joiner(matching_indexes)
+                index = joiner(matching_indexes)
+                joined_indexes[name] = index
+            else:
+                index = matching_indexes[0]
+
+        if dim in unlabeled_dim_sizes:
+            unlabeled_sizes = unlabeled_dim_sizes[dim]
+            labeled_size = index.size
+            if len(unlabeled_sizes | {labeled_size}) > 1:
+                raise ValueError(
+                    'arguments without labels along dimension %r cannot be '
+                    'aligned because they have different dimension size(s) %r '
+                    'than the size of the aligned dimension labels: %r'
+                    % (dim, unlabeled_sizes, labeled_size))
 
     for dim in unlabeled_dim_sizes:
-        if dim not in joined_indexes:
-            sizes = set(unlabeled_dim_sizes[dim])
+        if dim not in all_indexes:
+            sizes = unlabeled_dim_sizes[dim]
             if len(sizes) > 1:
                 raise ValueError(
-                    'dimension %r without indexes cannot be aligned because '
-                    'it has different sizes: %r' % (dim, sizes))
+                    'arguments without labels along dimension %r cannot be '
+                    'aligned because they have different dimension sizes: %r'
+                    % (dim, sizes))
 
     result = []
     for obj in objects:
-        valid_indexers = dict((k, v) for k, v in joined_indexes.items()
-                              if k in obj.dims)
+        valid_indexers = {k: v for k, v in joined_indexes.items()
+                          if k in obj.dims}
         result.append(obj.reindex(copy=copy, **valid_indexers))
 
     return tuple(result)
@@ -298,11 +318,6 @@ def broadcast(*args, **kwargs):
         The same data as the input arrays, but with additional dimensions
         inserted so that all data arrays have the same dimensions and shape.
 
-    Raises
-    ------
-    ValueError
-        If indexes on the different objects are not aligned.
-
     Examples
     --------
 

xarray/core/coordinates.py

@@ -11,13 +11,7 @@
 
 class AbstractCoordinates(Mapping, formatting.ReprMixin):
     def __getitem__(self, key):
-        if (key in self._names or
-            (isinstance(key, basestring) and
-             key.split('.')[0] in self._names)):
-            # allow indexing current coordinates or components
-            return self._data[key]
-        else:
-            raise KeyError(key)
+        raise NotImplementedError
 
     def __setitem__(self, key, value):
         self.update({key: value})
@@ -143,6 +137,11 @@ def variables(self):
                                   for k, v in self._data.variables.items()
                                   if k in self._names))
 
+    def __getitem__(self, key):
+        if key in self._data.data_vars:
+            raise KeyError(key)
+        return self._data[key]
+
     def to_dataset(self):
         """Convert these coordinates into a new Dataset
         """
@@ -188,6 +187,9 @@ def __init__(self, dataarray):
     def _names(self):
         return set(self._data._coords)
 
+    def __getitem__(self, key):
+        return self._data._getitem_coord(key)
+
     def _update_coords(self, coords):
         from .dataset import calculate_dimensions
 

xarray/core/dataarray.py

@@ -439,17 +439,21 @@ def _level_coords(self):
                     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
+    def _getitem_coord(self, key):
+        from .dataset import _get_virtual_variable
 
-            try:
-                var = self._coords[key]
-            except KeyError:
-                _, key, var = _get_virtual_variable(
-                    self._coords, key, self._level_coords)
+        try:
+            var = self._coords[key]
+        except KeyError:
+            dim_sizes = dict(zip(self.dims, self.shape))
+            _, key, var = _get_virtual_variable(
+                self._coords, key, self._level_coords, dim_sizes)
+
+        return self._replace_maybe_drop_dims(var, name=key)
 
-            return self._replace_maybe_drop_dims(var, name=key)
+    def __getitem__(self, key):
+        if isinstance(key, basestring):
+            return self._getitem_coord(key)
         else:
             # orthogonal array indexing
             return self.isel(**self._item_key_to_dict(key))

xarray/core/dataset.py

@@ -33,10 +33,19 @@
                              'quarter']
 
 
-def _get_virtual_variable(variables, key, level_vars={}):
+def _get_virtual_variable(variables, key, level_vars=None, dim_sizes=None):
     """Get a virtual variable (e.g., 'time.year' or a MultiIndex level)
     from a dict of xarray.Variable objects (if possible)
     """
+    if level_vars is None:
+        level_vars = {}
+    if dim_sizes is None:
+        dim_sizes = {}
+
+    if key in dim_sizes:
+        variable = Variable((key,), np.arange(dim_sizes[key]))
+        return key, key, variable
+
     if not isinstance(key, basestring):
         raise KeyError(key)
 
@@ -452,9 +461,9 @@ def _copy_listed(self, names):
                 variables[name] = self._variables[name]
             except KeyError:
                 ref_name, var_name, var = _get_virtual_variable(
-                    self._variables, name, self._level_coords)
+                    self._variables, name, self._level_coords, self.dims)
                 variables[var_name] = var
-                if ref_name in self._coord_names:
+                if ref_name in self._coord_names or ref_name in self.dims:
                     coord_names.add(var_name)
 
         return self._subset_with_all_valid_coords(variables, coord_names,
@@ -469,7 +478,7 @@ def _construct_dataarray(self, name):
             variable = self._variables[name]
         except KeyError:
             _, name, variable = _get_virtual_variable(
-                self._variables, name, self._level_coords)
+                self._variables, name, self._level_coords, self.dims)
 
         coords = OrderedDict()
         needed_dims = set(variable.dims)

xarray/core/indexing.py

@@ -266,14 +266,22 @@ def remap_label_indexers(data_obj, indexers, method=None, tolerance=None):
     if method is not None and not isinstance(method, str):
         raise TypeError('``method`` must be a string')
 
-    pos_indexers, new_indexes = {}, {}
-    for dim, label in iteritems(get_dim_indexers(data_obj, indexers)):
-        index = data_obj[dim].to_index()
-        idxr, new_idx = convert_label_indexer(index, label,
-                                              dim, method, tolerance)
-        pos_indexers[dim] = idxr
-        if new_idx is not None:
-            new_indexes[dim] = new_idx
+    pos_indexers = {}
+    new_indexes = {}
+
+    dim_indexers = get_dim_indexers(data_obj, indexers)
+    for dim, label in iteritems(dim_indexers):
+        try:
+            index_coord = data_obj[dim]
+        except KeyError:
+            pos_indexers[dim] = label
+        else:
+            index = index_coord.to_index()
+            idxr, new_idx = convert_label_indexer(index, label,
+                                                  dim, method, tolerance)
+            pos_indexers[dim] = idxr
+            if new_idx is not None:
+                new_indexes[dim] = new_idx
 
     return pos_indexers, new_indexes
 

xarray/test/test_dataarray.py

@@ -11,8 +11,8 @@
 from xarray.core.pycompat import iteritems, OrderedDict
 from xarray.core.common import _full_like
 
-from xarray.test import (TestCase, ReturnItem, source_ndarray, unittest, requires_dask,
-               requires_bottleneck)
+from xarray.test import (TestCase, ReturnItem, source_ndarray, unittest,
+                         requires_dask, requires_bottleneck)
 
 
 class TestDataArray(TestCase):
@@ -609,7 +609,13 @@ def test_sel(lab_indexer, pos_indexer, replaced_idx=False,
         self.assertDataArrayIdentical(mdata.sel(x={'one': 'a', 'two': 1}),
                                       mdata.sel(one='a', two=1))
 
-    def test_time_components(self):
+    def test_virtual_default_coords(self):
+        array = DataArray(np.zeros((5,)), dims='x')
+        expected = DataArray(range(5), dims='x', name='x')
+        self.assertDataArrayIdentical(expected, array['x'])
+        self.assertDataArrayIdentical(expected, array.coords['x'])
+
+    def test_virtual_time_components(self):
         dates = pd.date_range('2000-01-01', periods=10)
         da = DataArray(np.arange(1, 11), [('time', dates)])
 
@@ -743,7 +749,9 @@ def test_coords_alignment(self):
         rhs = DataArray([2, 3, 4], [('x', [1, 2, 3])])
         lhs.coords['rhs'] = rhs
 
-        expected = DataArray([1, 2, 3], coords={'rhs': ('x', [np.nan, 2, 3])},
+        expected = DataArray([1, 2, 3],
+                             coords={'rhs': ('x', [np.nan, 2, 3]),
+                                     'x': [0, 1, 2]},
                              dims='x')
         self.assertDataArrayIdentical(lhs, expected)
 
@@ -755,6 +763,12 @@ def test_coords_replacement_alignment(self):
         expected = DataArray([0, 1, 2], coords=[('abc', [1, 2, 3])])
         self.assertDataArrayIdentical(arr, expected)
 
+    def test_coords_non_string(self):
+        arr = DataArray(0, coords={1: 2})
+        actual = arr.coords[1]
+        expected = DataArray(2, coords={1: 2}, name=1)
+        self.assertDataArrayIdentical(actual, expected)
+
     def test_reindex(self):
         foo = self.dv
         bar = self.dv[:2, :2]
@@ -1649,10 +1663,11 @@ def test_resample_upsampling(self):
             self.assertDataArrayIdentical(expected, actual)
 
     def test_align(self):
-        self.ds['x'] = ('x', np.array(list('abcdefghij')))
-        dv1, dv2 = align(self.dv, self.dv[:5], join='inner')
-        self.assertDataArrayIdentical(dv1, self.dv[:5])
-        self.assertDataArrayIdentical(dv2, self.dv[:5])
+        array = DataArray(np.random.random((6, 8)),
+                          coords={'x': list('abcdef')}, dims=['x', 'y'])
+        array1, array2 = align(array, array[:5], join='inner')
+        self.assertDataArrayIdentical(array1, array[:5])
+        self.assertDataArrayIdentical(array2, array[:5])
 
     def test_align_dtype(self):
         # regression test for #264
@@ -1723,6 +1738,15 @@ def test_align_indexes(self):
                                 coords=[('a', [-2, 7, 10, -1])])
         self.assertDataArrayIdentical(expected_x2, x2)
 
+    def test_align_without_indexes_errors(self):
+        with self.assertRaisesRegexp(ValueError, 'cannot be aligned'):
+            align(DataArray([1, 2, 3], dims=['x']),
+                  DataArray([1, 2], dims=['x']))
+
+        with self.assertRaisesRegexp(ValueError, 'cannot be aligned'):
+            align(DataArray([1, 2, 3], dims=['x']),
+                  DataArray([1, 2], coords=[('x', [0, 1])]))
+
     def test_broadcast_arrays(self):
         x = DataArray([1, 2], coords=[('a', [-1, -2])], name='x')
         y = DataArray([1, 2], coords=[('b', [3, 4])], name='y')

xarray/test/test_dataset.py

@@ -33,7 +33,6 @@ def create_test_data(seed=None):
 
     obj = Dataset()
     obj['time'] = ('time', pd.date_range('2000-01-01', periods=20))
-    obj['dim1'] = ('dim1', np.arange(_dims['dim1'], dtype='int64'))
     obj['dim2'] = ('dim2', 0.5 * np.arange(_dims['dim2']))
     obj['dim3'] = ('dim3', list('abcdefghij'))
     for v, dims in sorted(_vars.items()):
@@ -70,7 +69,6 @@ def test_repr(self):
         Dimensions:  (dim1: 8, dim2: 9, dim3: 10, time: 20)
         Coordinates:
           * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 ...
-          * dim1     (dim1) int64 0 1 2 3 4 5 6 7
           * dim2     (dim2) float64 0.0 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0
           * dim3     (dim3) %s 'a' 'b' 'c' 'd' 'e' 'f' 'g' 'h' 'i' 'j'
             numbers  (dim3) int64 0 1 2 0 0 1 1 2 2 3
@@ -566,10 +564,12 @@ def test_coords_to_dataset(self):
         self.assertDatasetIdentical(expected, actual)
 
     def test_coords_merge(self):
-        orig_coords = Dataset(coords={'a': ('x', [1, 2])}).coords
-        other_coords = Dataset(coords={'b': ('x', ['a', 'b'])}).coords
+        orig_coords = Dataset(coords={'a': ('x', [1, 2]), 'x': [0, 1]}).coords
+        other_coords = Dataset(coords={'b': ('x', ['a', 'b']),
+                                       'x': [0, 1]}).coords
         expected = Dataset(coords={'a': ('x', [1, 2]),
-                                   'b': ('x', ['a', 'b'])})
+                                   'b': ('x', ['a', 'b']),
+                                   'x': [0, 1]})
         actual = orig_coords.merge(other_coords)
         self.assertDatasetIdentical(expected, actual)
         actual = other_coords.merge(orig_coords)
@@ -1495,7 +1495,17 @@ def test_getitem_hashable(self):
         with self.assertRaisesRegexp(KeyError, "('var1', 'var2')"):
             data[('var1', 'var2')]
 
-    def test_virtual_variables(self):
+    def test_virtual_variables_default_coords(self):
+        dataset = Dataset({'foo': ('x', range(10))})
+        expected = DataArray(range(10), dims='x', name='x')
+        actual = dataset['x']
+        self.assertDataArrayIdentical(expected, actual)
+
+        actual = dataset[['x', 'foo']]
+        expected = dataset.assign_coords(x=range(10))
+        self.assertDatasetIdentical(expected, actual)
+
+    def test_virtual_variables_time(self):
         # access virtual variables
         data = create_test_data()
         expected = DataArray(1 + np.arange(20), coords=[data['time']],