added tests for indexing and selection

Author: Benoit Bovy

xarray/test/test_dataarray.py

@@ -486,7 +486,8 @@ def test_loc_single_boolean(self):
         self.assertEqual(data.loc[False], 1)
 
     def test_multiindex(self):
-        idx = pd.MultiIndex.from_product([list('abc'), [0, 1]])
+        idx = pd.MultiIndex.from_product([list('abc'), [0, 1]],
+                                         names=('one', 'two'))
         data = DataArray(range(6), [('x', idx)])
 
         self.assertDataArrayIdentical(data.sel(x=('a', 0)), data.isel(x=0))
@@ -495,6 +496,20 @@ def test_multiindex(self):
         self.assertDataArrayIdentical(data.sel(x=[('a', 0), ('c', 1)]),
                                       data.isel(x=[0, -1]))
         self.assertDataArrayIdentical(data.sel(x='a'), data.isel(x=slice(2)))
+        self.assertVariableNotEqual(data.sel(x={'one': slice(None)}), data)
+        self.assertDataArrayIdentical(data.isel(x=[0]),
+                                      data.sel(x={'one': 'a', 'two': 0}))
+        self.assertDataArrayIdentical(data.isel(x=[0, 1]), data.sel(x='a'))
+        self.assertVariableIdentical(
+            data.sel(x={'one': 'a'}),
+            data.unstack('x').sel(one='a').dropna('two')
+        )
+
+        self.assertDataArrayIdentical(data.loc['a'], data[:2])
+        self.assertDataArrayIdentical(data.loc[{'one': 'a', 'two': 0}, ...],
+                                       data[[0]])
+        self.assertDataArrayIdentical(data.loc[{'one': 'a'}, ...],
+                                      data.sel(x={'one': 'a'}))
 
     def test_time_components(self):
         dates = pd.date_range('2000-01-01', periods=10)
@@ -1792,29 +1807,29 @@ def test_full_like(self):
         actual = _full_like(DataArray([1, 2, 3]), fill_value=np.nan)
         self.assertEqual(actual.dtype, np.float)
         np.testing.assert_equal(actual.values, np.nan)
-    
+
     def test_dot(self):
         x = np.linspace(-3, 3, 6)
         y = np.linspace(-3, 3, 5)
-        z = range(4) 
+        z = range(4)
         da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4))
         da = DataArray(da_vals, coords=[x, y, z], dims=['x', 'y', 'z'])
-        
+
         dm_vals = range(4)
         dm = DataArray(dm_vals, coords=[z], dims=['z'])
-        
+
         # nd dot 1d
         actual = da.dot(dm)
         expected_vals = np.tensordot(da_vals, dm_vals, [2, 0])
         expected = DataArray(expected_vals, coords=[x, y], dims=['x', 'y'])
         self.assertDataArrayEqual(expected, actual)
-        
+
         # all shared dims
         actual = da.dot(da)
         expected_vals = np.tensordot(da_vals, da_vals, axes=([0, 1, 2], [0, 1, 2]))
         expected = DataArray(expected_vals)
         self.assertDataArrayEqual(expected, actual)
-        
+
         # multiple shared dims
         dm_vals = np.arange(20 * 5 * 4).reshape((20, 5, 4))
         j = np.linspace(-3, 3, 20)
@@ -1823,7 +1838,7 @@ def test_dot(self):
         expected_vals = np.tensordot(da_vals, dm_vals, axes=([1, 2], [1, 2]))
         expected = DataArray(expected_vals, coords=[x, j], dims=['x', 'j'])
         self.assertDataArrayEqual(expected, actual)
-        
+
         with self.assertRaises(NotImplementedError):
             da.dot(dm.to_dataset(name='dm'))
         with self.assertRaises(TypeError):

xarray/test/test_dataset.py

@@ -840,6 +840,31 @@ def test_loc(self):
         with self.assertRaises(TypeError):
             data.loc[dict(dim3='a')] = 0
 
+    def test_multiindex(self):
+        idx = pd.MultiIndex.from_product([list('abc'), [0, 1]],
+                                         names=('one', 'two'))
+        data = Dataset(data_vars={'var': ('x', range(6))}, coords={'x': idx})
+
+        self.assertDatasetIdentical(data.sel(x=('a', 0)), data.isel(x=0))
+        self.assertDatasetIdentical(data.sel(x=('c', 1)), data.isel(x=-1))
+        self.assertDatasetIdentical(data.sel(x=[('a', 0)]), data.isel(x=[0]))
+        self.assertDatasetIdentical(data.sel(x=[('a', 0), ('c', 1)]),
+                                      data.isel(x=[0, -1]))
+        self.assertDatasetIdentical(data.sel(x='a'), data.isel(x=slice(2)))
+        self.assertVariableNotEqual(data.sel(x={'one': slice(None)})['var'],
+                                    data['var'])
+        self.assertDatasetIdentical(data.isel(x=[0]),
+                                      data.sel(x={'one': 'a', 'two': 0}))
+        self.assertDatasetIdentical(data.isel(x=[0, 1]), data.sel(x='a'))
+        self.assertVariableIdentical(
+            data.sel(x={'one': 'a'})['var'],
+            data.unstack('x').sel(one='a').dropna('two')['var']
+        )
+
+        self.assertDatasetIdentical(data.loc[{'x': 'a'}], data.sel(x='a'))
+        self.assertDatasetIdentical(data.loc[{'x': {'one': 'a', 'two': 0}}],
+                                    data.sel(x={'one': 'a', 'two': 0}))
+
     def test_reindex_like(self):
         data = create_test_data()
         data['letters'] = ('dim3', 10 * ['a'])

xarray/test/test_indexing.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 
-from xarray import Dataset, Variable
+from xarray import Dataset, DataArray, Variable
 from xarray.core import indexing
 from . import TestCase, ReturnItem
 
@@ -85,8 +85,18 @@ def test_convert_label_indexer(self):
             indexing.convert_label_indexer(index, [0])
         with self.assertRaises(KeyError):
             indexing.convert_label_indexer(index, 0)
+        with self.assertRaisesRegexp(ValueError, 'does not have a MultiIndex'):
+            indexing.convert_label_indexer(index, {'one': 0})
+
+        mindex = pd.MultiIndex(levels=[[1, 2, 3], [4, 5, 6]],
+                               labels=[[0, 1, 2], [0, 1, 2]],
+                               names=['one', 'two'])
+        with self.assertRaises(KeyError):
+            indexing.convert_label_indexer(mindex, [0])
+        with self.assertRaises(KeyError):
+            indexing.convert_label_indexer(mindex, 0)
         with self.assertRaises(ValueError):
-            indexing.convert_label_indexer(index, {'somelevel': 0})
+            indexing.convert_label_indexer(index, {'three': 0})
 
     def test_convert_unsorted_datetime_index_raises(self):
         index = pd.to_datetime(['2001', '2000', '2002'])
@@ -105,6 +115,25 @@ def test_indexer(x):
         self.assertEqual({'x': 0}, test_indexer(1)[0])
         self.assertEqual({'x': 0}, test_indexer(np.int32(1))[0])
         self.assertEqual({'x': 0}, test_indexer(Variable([], 1))[0])
+        with self.assertRaisesRegexp(ValueError, 'does not have a MultiIndex'):
+            indexing.remap_label_indexers(data, {'x': {'level': 1}})
+
+        mindex = pd.MultiIndex(levels=[[1, 2, 3], [4, 5, 6]],
+                               labels=[[0, 1, 2], [0, 1, 2]],
+                               names=['one', 'two'])
+        s = pd.Series([7, 8, 9], index=mindex)
+        data = DataArray(s, dims='x')
+
+        pos, idx = indexing.remap_label_indexers(data, {'x': {'one': 1}})
+        self.assertArrayEqual(pos['x'], [True, False, False])
+        self.assertArrayEqual(idx['x'].values, [4])
+        self.assertEqual(idx['x'].name, 'two')
+
+        pos, idx = indexing.remap_label_indexers(
+            data, {'x': {'one': 1, 'two': 4}}
+        )
+        self.assertArrayEqual(pos['x'], [True, False, False])
+        self.assertEqual(len(idx), 0)
 
 
 class TestLazyArray(TestCase):