class TestMultiIndex(unittest.TestCase): def setUp(self): major_axis = Index(['foo', 'bar', 'baz', 'qux']) minor_axis = Index(['one', 'two']) major_labels = np.array([0, 0, 1, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) self.index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels], names=['first', 'second']) def test_constructor_single_level(self): single_level = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux']], labels=[[0, 1, 2, 3]], names=['first']) self.assert_(isinstance(single_level, Index)) self.assert_(not isinstance(single_level, MultiIndex)) self.assert_(single_level.name == 'first') single_level = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux']], labels=[[0, 1, 2, 3]]) self.assert_(single_level.name is None) def test_constructor_no_levels(self): self.assertRaises(Exception, MultiIndex, levels=[], labels=[]) def test_duplicate_names(self): self.index.names = ['foo', 'foo'] self.assertRaises(Exception, self.index._get_level_number, 'foo') def test_get_level_number_integer(self): self.index.names = [1, 0] self.assertEqual(self.index._get_level_number(1), 0) self.assertEqual(self.index._get_level_number(0), 1) self.assertRaises(Exception, self.index._get_level_number, 2) def test_from_arrays(self): arrays = [] for lev, lab in zip(self.index.levels, self.index.labels): arrays.append(np.asarray(lev).take(lab)) result = MultiIndex.from_arrays(arrays) self.assertEquals(list(result), list(self.index)) def test_append(self): result = self.index[:3].append(self.index[3:]) self.assert_(result.equals(self.index)) foos = [self.index[:1], self.index[1:3], self.index[3:]] result = foos[0].append(foos[1:]) self.assert_(result.equals(self.index)) # empty result = self.index.append([]) self.assert_(result.equals(self.index)) def test_get_level_values(self): result = self.index.get_level_values(0) expected = ['foo', 'foo', 'bar', 'baz', 'qux', 'qux'] self.assert_(np.array_equal(result, expected)) result = self.index.get_level_values('first') expected = self.index.get_level_values(0) self.assert_(np.array_equal(result, expected)) def test_nlevels(self): self.assertEquals(self.index.nlevels, 2) def test_iter(self): result = list(self.index) expected = [('foo', 'one'), ('foo', 'two'), ('bar', 'one'), ('baz', 'two'), ('qux', 'one'), ('qux', 'two')] self.assert_(result == expected) def test_pickle(self): pickled = pickle.dumps(self.index) unpickled = pickle.loads(pickled) self.assert_(self.index.equals(unpickled)) def test_legacy_pickle(self): if py3compat.PY3: raise nose.SkipTest import os def curpath(): pth, _ = os.path.split(os.path.abspath(__file__)) return pth ppath = os.path.join(curpath(), 'data/multiindex_v1.pickle') obj = pickle.load(open(ppath, 'r')) self.assert_(obj._is_legacy_format) obj2 = MultiIndex.from_tuples(obj.values) self.assert_(obj.equals(obj2)) res = obj.get_indexer(obj2[::-1]) exp = obj.get_indexer(obj[::-1]) exp2 = obj2.get_indexer(obj2[::-1]) assert_almost_equal(res, exp) assert_almost_equal(exp, exp2) def test_contains(self): self.assert_(('foo', 'two') in self.index) self.assert_(('bar', 'two') not in self.index) self.assert_(None not in self.index) def test_is_all_dates(self): self.assert_(not self.index.is_all_dates) def test_getitem(self): # scalar self.assertEquals(self.index[2], ('bar', 'one')) # slice result = self.index[2:5] expected = self.index[[2,3,4]] self.assert_(result.equals(expected)) # boolean result = self.index[[True, False, True, False, True, True]] result2 = self.index[np.array([True, False, True, False, True, True])] expected = self.index[[0, 2, 4, 5]] self.assert_(result.equals(expected)) self.assert_(result2.equals(expected)) def test_getitem_group_select(self): sorted_idx, _ = self.index.sortlevel(0) self.assertEquals(sorted_idx.get_loc('baz'), slice(3, 4)) self.assertEquals(sorted_idx.get_loc('foo'), slice(0, 2)) def test_get_loc(self): self.assert_(self.index.get_loc(('foo', 'two')) == 1) self.assert_(self.index.get_loc(('baz', 'two')) == 3) self.assertRaises(KeyError, self.index.get_loc, ('bar', 'two')) self.assertRaises(KeyError, self.index.get_loc, 'quux') # 3 levels index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])]) self.assertRaises(KeyError, index.get_loc, (1, 1)) self.assert_(index.get_loc((2, 0)) == slice(3, 5)) def test_get_loc_duplicates(self): index = Index([2, 2, 2, 2]) self.assertRaises(Exception, index.get_loc, 2) def test_get_loc_level(self): index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])]) loc, new_index = index.get_loc_level((0, 1)) expected = slice(1, 2) exp_index = index[expected].droplevel(0).droplevel(0) self.assertEqual(loc, expected) self.assert_(new_index.equals(exp_index)) loc, new_index = index.get_loc_level((0, 1, 0)) expected = 1 self.assertEqual(loc, expected) self.assert_(new_index is None) self.assertRaises(KeyError, index.get_loc_level, (2, 2)) index = MultiIndex(levels=[[2000], range(4)], labels=[np.array([0, 0, 0, 0]), np.array([0, 1, 2, 3])]) result, new_index = index.get_loc_level((2000, slice(None, None))) expected = slice(None, None) self.assertEqual(result, expected) self.assert_(new_index.equals(index.droplevel(0))) def test_slice_locs(self): df = tm.makeTimeDataFrame() stacked = df.stack() idx = stacked.index slob = slice(*idx.slice_locs(df.index[5], df.index[15])) sliced = stacked[slob] expected = df[5:16].stack() tm.assert_almost_equal(sliced.values, expected.values) slob = slice(*idx.slice_locs(df.index[5] + timedelta(seconds=30), df.index[15] - timedelta(seconds=30))) sliced = stacked[slob] expected = df[6:15].stack() tm.assert_almost_equal(sliced.values, expected.values) def test_slice_locs_not_sorted(self): index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])]) self.assertRaises(Exception, index.slice_locs, (1, 0, 1), (2, 1, 0)) # works sorted_index, _ = index.sortlevel(0) result = sorted_index.slice_locs((1, 0, 1), (2, 1, 0)) def test_slice_locs_partial(self): sorted_idx, _ = self.index.sortlevel(0) result = sorted_idx.slice_locs(('foo', 'two'), ('qux', 'one')) self.assertEquals(result, (1, 5)) result = sorted_idx.slice_locs(None, ('qux', 'one')) self.assertEquals(result, (0, 5)) result = sorted_idx.slice_locs(('foo', 'two'), None) self.assertEquals(result, (1, len(sorted_idx))) result = sorted_idx.slice_locs('bar', 'baz') self.assertEquals(result, (2, 4)) def test_slice_locs_not_contained(self): # some searchsorted action index = MultiIndex(levels=[[0, 2, 4, 6], [0, 2, 4]], labels=[[0, 0, 0, 1, 1, 2, 3, 3, 3], [0, 1, 2, 1, 2, 2, 0, 1, 2]], sortorder=0) result = index.slice_locs((1, 0), (5, 2)) self.assertEquals(result, (3, 6)) result = index.slice_locs(1, 5) self.assertEquals(result, (3, 6)) result = index.slice_locs((2, 2), (5, 2)) self.assertEquals(result, (3, 6)) result = index.slice_locs(2, 5) self.assertEquals(result, (3, 6)) result = index.slice_locs((1, 0), (6, 3)) self.assertEquals(result, (3, 8)) result = index.slice_locs(-1, 10) self.assertEquals(result, (0, len(index))) def test_consistency(self): # need to construct an overflow major_axis = range(70000) minor_axis = range(10) major_labels = np.arange(70000) minor_labels = np.repeat(range(10), 7000) # the fact that is works means it's consistent index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) # inconsistent major_labels = np.array([0, 0, 1, 1, 1, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assertRaises(Exception, getattr, index, 'indexMap') def test_truncate(self): major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) result = index.truncate(before=1) self.assert_('foo' not in result.levels[0]) self.assert_(1 in result.levels[0]) result = index.truncate(after=1) self.assert_(2 not in result.levels[0]) self.assert_(1 in result.levels[0]) result = index.truncate(before=1, after=2) self.assertEqual(len(result.levels[0]), 2) # after < before self.assertRaises(ValueError, index.truncate, 3, 1) def test_get_indexer(self): major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) idx1 = index[:5] idx2 = index[[1,3,5]] r1 = idx1.get_indexer(idx2) assert_almost_equal(r1, [1, 3, -1]) r1 = idx2.get_indexer(idx1, method='pad') assert_almost_equal(r1, [-1, 0, 0, 1, 1]) rffill1 = idx2.get_indexer(idx1, method='ffill') assert_almost_equal(r1, rffill1) r1 = idx2.get_indexer(idx1, method='backfill') assert_almost_equal(r1, [0, 0, 1, 1, 2]) rbfill1 = idx2.get_indexer(idx1, method='bfill') assert_almost_equal(r1, rbfill1) # pass non-MultiIndex r1 = idx1.get_indexer(idx2.get_tuple_index()) rexp1 = idx1.get_indexer(idx2) assert_almost_equal(r1, rexp1) r1 = idx1.get_indexer([1,2,3]) self.assert_( (r1 == [-1, -1, -1]).all() ) # self.assertRaises(Exception, idx1.get_indexer, # list(list(zip(*idx2.get_tuple_index()))[0])) def test_format(self): self.index.format() self.index[:0].format() def test_format_integer_names(self): index = MultiIndex(levels=[[0, 1], [0, 1]], labels=[[0, 0, 1, 1], [0, 1, 0, 1]], names=[0, 1]) index.format(names=True) def test_bounds(self): self.index._bounds def test_equals(self): self.assert_(self.index.equals(self.index)) self.assert_(self.index.equal_levels(self.index)) self.assert_(not self.index.equals(self.index[:-1])) self.assert_(self.index.equals(self.index.get_tuple_index())) # different number of levels index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])]) index2 = MultiIndex(levels=index.levels[:-1], labels=index.labels[:-1]) self.assert_(not index.equals(index2)) self.assert_(not index.equal_levels(index2)) # levels are different major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 2, 3]) minor_labels = np.array([0, 1, 0, 0, 1, 0]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assert_(not self.index.equals(index)) self.assert_(not self.index.equal_levels(index)) # some of the labels are different major_axis = Index(['foo', 'bar', 'baz', 'qux']) minor_axis = Index(['one', 'two']) major_labels = np.array([0, 0, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assert_(not self.index.equals(index)) def test_union(self): piece1 = self.index[:5][::-1] piece2 = self.index[3:] the_union = piece1 | piece2 tups = sorted(self.index.get_tuple_index()) expected = MultiIndex.from_tuples(tups) self.assert_(the_union.equals(expected)) # corner case, pass self or empty thing: the_union = self.index.union(self.index) self.assert_(the_union is self.index) the_union = self.index.union(self.index[:0]) self.assert_(the_union is self.index) # won't work in python 3 # tuples = self.index.get_tuple_index() # result = self.index[:4] | tuples[4:] # self.assert_(result.equals(tuples)) # not valid for python 3 # def test_union_with_regular_index(self): # other = Index(['A', 'B', 'C']) # result = other.union(self.index) # self.assert_(('foo', 'one') in result) # self.assert_('B' in result) # result2 = self.index.union(other) # self.assert_(result.equals(result2)) def test_intersection(self): piece1 = self.index[:5][::-1] piece2 = self.index[3:] the_int = piece1 & piece2 tups = sorted(self.index[3:5].get_tuple_index()) expected = MultiIndex.from_tuples(tups) self.assert_(the_int.equals(expected)) # corner case, pass self the_int = self.index.intersection(self.index) self.assert_(the_int is self.index) # empty intersection: disjoint empty = self.index[:2] & self.index[2:] expected = self.index[:0] self.assert_(empty.equals(expected)) # can't do in python 3 # tuples = self.index.get_tuple_index() # result = self.index & tuples # self.assert_(result.equals(tuples)) def test_diff(self): first = self.index result = first - self.index[-3:] expected = MultiIndex.from_tuples(sorted(self.index[:-3].values), sortorder=0, names=self.index.names) self.assert_(isinstance(result, MultiIndex)) self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # empty difference: reflexive result = self.index - self.index expected = self.index[:0] self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # empty difference: superset result = self.index[-3:] - self.index expected = self.index[:0] self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # empty difference: degenerate result = self.index[:0] - self.index expected = self.index[:0] self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # names not the same chunklet = self.index[-3:] chunklet.names = ['foo', 'baz'] result = first - chunklet self.assertEqual(result.names, [None, None]) # empty, but non-equal result = self.index - self.index.sortlevel(1)[0] self.assert_(len(result) == 0) # raise Exception called with non-MultiIndex self.assertRaises(Exception, first.diff, first.get_tuple_index()) def test_from_tuples(self): self.assertRaises(Exception, MultiIndex.from_tuples, []) def test_argsort(self): result = self.index.argsort() expected = self.index.get_tuple_index().argsort() self.assert_(np.array_equal(result, expected)) def test_sortlevel(self): import random tuples = list(self.index) random.shuffle(tuples) index = MultiIndex.from_tuples(tuples) sorted_idx, _ = index.sortlevel(0) expected = MultiIndex.from_tuples(sorted(tuples)) self.assert_(sorted_idx.equals(expected)) sorted_idx, _ = index.sortlevel(0, ascending=False) self.assert_(sorted_idx.equals(expected[::-1])) sorted_idx, _ = index.sortlevel(1) by1 = sorted(tuples, key=lambda x: (x[1], x[0])) expected = MultiIndex.from_tuples(by1) self.assert_(sorted_idx.equals(expected)) sorted_idx, _ = index.sortlevel(1, ascending=False) self.assert_(sorted_idx.equals(expected[::-1])) def test_dims(self): pass def test_drop(self): dropped = self.index.drop([('foo', 'two'), ('qux', 'one')]) index = MultiIndex.from_tuples([('foo', 'two'), ('qux', 'one')]) dropped2 = self.index.drop(index) expected = self.index[[0, 2, 3, 5]] self.assert_(dropped.equals(expected)) self.assert_(dropped2.equals(expected)) dropped = self.index.drop(['bar']) expected = self.index[[0, 1, 3, 4, 5]] self.assert_(dropped.equals(expected)) index = MultiIndex.from_tuples([('bar', 'two')]) self.assertRaises(Exception, self.index.drop, [('bar', 'two')]) self.assertRaises(Exception, self.index.drop, index) # mixed partial / full drop dropped = self.index.drop(['foo', ('qux', 'one')]) expected = self.index[[2, 3, 5]] self.assert_(dropped.equals(expected)) def test_droplevel_with_names(self): index = self.index[self.index.get_loc('foo')] dropped = index.droplevel(0) self.assertEqual(dropped.name, 'second') index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])], names=['one', 'two', 'three']) dropped = index.droplevel(0) self.assertEqual(dropped.names, ['two', 'three']) dropped = index.droplevel('two') expected = index.droplevel(1) self.assert_(dropped.equals(expected)) def test_droplevel_multiple(self): index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])], names=['one', 'two', 'three']) dropped = index[:2].droplevel(['three', 'one']) expected = index[:2].droplevel(2).droplevel(0) self.assert_(dropped.equals(expected)) def test_insert(self): # key contained in all levels new_index = self.index.insert(0, ('bar', 'two')) self.assert_(new_index.equal_levels(self.index)) self.assert_(new_index[0] == ('bar', 'two')) # key not contained in all levels new_index = self.index.insert(0, ('abc', 'three')) self.assert_(np.array_equal(new_index.levels[0], list(self.index.levels[0]) + ['abc'])) self.assert_(np.array_equal(new_index.levels[1], list(self.index.levels[1]) + ['three'])) self.assert_(new_index[0] == ('abc', 'three')) # key wrong length self.assertRaises(Exception, self.index.insert, 0, ('foo2',)) def test_take_preserve_name(self): taken = self.index.take([3,0,1]) self.assertEqual(taken.names, self.index.names) def test_join_level(self): def _check_how(other, how): join_index, lidx, ridx = other.join(self.index, how=how, level='second', return_indexers=True) exp_level = other.join(self.index.levels[1], how=how) self.assert_(join_index.levels[0].equals(self.index.levels[0])) self.assert_(join_index.levels[1].equals(exp_level)) # pare down levels mask = np.array([x[1] in exp_level for x in self.index], dtype=bool) exp_values = self.index.values[mask] self.assert_(np.array_equal(join_index.values, exp_values)) if how in ('outer', 'inner'): join_index2, ridx2, lidx2 = \ self.index.join(other, how=how, level='second', return_indexers=True) self.assert_(join_index.equals(join_index2)) self.assert_(np.array_equal(lidx, lidx2)) self.assert_(np.array_equal(ridx, ridx2)) self.assert_(np.array_equal(join_index2.values, exp_values)) def _check_all(other): _check_how(other, 'outer') _check_how(other, 'inner') _check_how(other, 'left') _check_how(other, 'right') _check_all(Index(['three', 'one', 'two'])) _check_all(Index(['one'])) _check_all(Index(['one', 'three'])) # some corner cases idx = Index(['three', 'one', 'two']) result = idx.join(self.index, level='second') self.assert_(isinstance(result, MultiIndex)) self.assertRaises(Exception, self.index.join, self.index, level=1) def test_reindex(self): result, indexer = self.index.reindex(list(self.index[:4])) self.assert_(isinstance(result, MultiIndex)) result, indexer = self.index.reindex(list(self.index)) self.assert_(isinstance(result, MultiIndex)) self.assert_(indexer is None) def test_reindex_level(self): idx = Index(['one']) target, indexer = self.index.reindex(idx, level='second') target2, indexer2 = idx.reindex(self.index, idx, level='second') exp_index = self.index.join(idx, level='second', how='left') self.assert_(target.equals(exp_index)) self.assert_(target2.equals(exp_index)) def test_has_duplicates(self): self.assert_(not self.index.has_duplicates) self.assert_(self.index.append(self.index).has_duplicates) index = MultiIndex(levels=[[0, 1], [0, 1, 2]], labels=[[0, 0, 0, 0, 1, 1, 1], [0, 1, 2, 0, 0, 1, 2]]) self.assert_(index.has_duplicates)
class TestMultiIndex(unittest.TestCase): def setUp(self): major_axis = Index(['foo', 'bar', 'baz', 'qux']) minor_axis = Index(['one', 'two']) major_labels = np.array([0, 0, 1, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) self.index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) def test_from_arrays(self): arrays = [] for lev, lab in zip(self.index.levels, self.index.labels): arrays.append(np.asarray(lev).take(lab)) result = MultiIndex.from_arrays(arrays) self.assertEquals(list(result), list(self.index)) def test_nlevels(self): self.assertEquals(self.index.nlevels, 2) def test_iter(self): result = list(self.index) expected = [('foo', 'one'), ('foo', 'two'), ('bar', 'one'), ('baz', 'two'), ('qux', 'one'), ('qux', 'two')] self.assert_(result == expected) def test_pickle(self): import pickle pickled = pickle.dumps(self.index) unpickled = pickle.loads(pickled) self.assert_(self.index.equals(unpickled)) def test_contains(self): self.assert_(('foo', 'two') in self.index) self.assert_(('bar', 'two') not in self.index) self.assert_(None not in self.index) def test_is_all_dates(self): self.assert_(not self.index.is_all_dates()) def test_getitem(self): # scalar self.assertEquals(self.index[2], ('bar', 'one')) # slice result = self.index[2:5] expected = self.index[[2,3,4]] self.assert_(result.equals(expected)) # boolean result = self.index[[True, False, True, False, True, True]] result2 = self.index[np.array([True, False, True, False, True, True])] expected = self.index[[0, 2, 4, 5]] self.assert_(result.equals(expected)) self.assert_(result2.equals(expected)) def test_getitem_group_select(self): sorted_idx, _ = self.index.sortlevel(0) self.assertEquals(sorted_idx.get_loc('baz'), slice(3, 4)) self.assertEquals(sorted_idx.get_loc('foo'), slice(0, 2)) def test_get_loc(self): self.assert_(self.index.get_loc(('foo', 'two')) == 1) self.assert_(self.index.get_loc(('baz', 'two')) == 3) self.assertRaises(KeyError, self.index.get_loc, ('bar', 'two')) self.assertRaises(KeyError, self.index.get_loc, 'quux') # 3 levels index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])]) self.assertRaises(KeyError, index.get_loc, (1, 1)) self.assert_(index.get_loc((2, 0)) == slice(3, 5)) def test_slice_locs(self): df = tm.makeTimeDataFrame() stacked = df.stack() idx = stacked.index slob = slice(*idx.slice_locs(df.index[5], df.index[15])) sliced = stacked[slob] expected = df[5:16].stack() tm.assert_almost_equal(sliced.values, expected.values) slob = slice(*idx.slice_locs(df.index[5] + timedelta(seconds=30), df.index[15] - timedelta(seconds=30))) sliced = stacked[slob] expected = df[6:15].stack() tm.assert_almost_equal(sliced.values, expected.values) def test_slice_locs_not_sorted(self): index = MultiIndex(levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0])]) self.assertRaises(Exception, index.slice_locs, (1, 0, 1), (2, 1, 0)) # works sorted_index, _ = index.sortlevel(0) result = sorted_index.slice_locs((1, 0, 1), (2, 1, 0)) def test_slice_locs_partial(self): sorted_idx, _ = self.index.sortlevel(0) result = sorted_idx.slice_locs(('foo', 'two'), ('qux', 'one')) self.assertEquals(result, (1, 5)) result = sorted_idx.slice_locs(None, ('qux', 'one')) self.assertEquals(result, (0, 5)) result = sorted_idx.slice_locs(('foo', 'two'), None) self.assertEquals(result, (1, len(sorted_idx))) result = sorted_idx.slice_locs('bar', 'baz') self.assertEquals(result, (2, 4)) def test_slice_locs_not_contained(self): # some searchsorted action index = MultiIndex(levels=[[0, 2, 4, 6], [0, 2, 4]], labels=[[0, 0, 0, 1, 1, 2, 3, 3, 3], [0, 1, 2, 1, 2, 2, 0, 1, 2]], sortorder=0) result = index.slice_locs((1, 0), (5, 2)) self.assertEquals(result, (3, 6)) result = index.slice_locs(1, 5) self.assertEquals(result, (3, 6)) result = index.slice_locs((2, 2), (5, 2)) self.assertEquals(result, (3, 6)) result = index.slice_locs(2, 5) self.assertEquals(result, (3, 6)) result = index.slice_locs((1, 0), (6, 3)) self.assertEquals(result, (3, 8)) result = index.slice_locs(-1, 10) self.assertEquals(result, (0, len(index))) def test_consistency(self): # need to construct an overflow major_axis = range(70000) minor_axis = range(10) major_labels = np.arange(70000) minor_labels = np.repeat(range(10), 7000) # the fact that is works means it's consistent index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) # inconsistent major_labels = np.array([0, 0, 1, 1, 1, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assertRaises(Exception, getattr, index, 'indexMap') def test_truncate(self): major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) result = index.truncate(before=1) self.assert_('foo' not in result.levels[0]) self.assert_(1 in result.levels[0]) result = index.truncate(after=1) self.assert_(2 not in result.levels[0]) self.assert_(1 in result.levels[0]) result = index.truncate(before=1, after=2) self.assertEqual(len(result.levels[0]), 2) def test_get_indexer(self): major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) idx1 = index[:5] idx2 = index[[1,3,5]] r1, r2 = idx1.get_indexer(idx2) assert_almost_equal(r1, [1, 3, -1]) assert_almost_equal(r2, [True, True, False]) r1, r2 = idx2.get_indexer(idx1, method='pad') assert_almost_equal(r1, [-1, 0, 0, 1, 1]) assert_almost_equal(r2, [False, True, True, True, True]) rffill1, rffill2 = idx2.get_indexer(idx1, method='ffill') assert_almost_equal(r1, rffill1) assert_almost_equal(r2, rffill2) r1, r2 = idx2.get_indexer(idx1, method='backfill') assert_almost_equal(r1, [0, 0, 1, 1, 2]) assert_almost_equal(r2, [True, True, True, True, True]) rbfill1, rbfill2 = idx2.get_indexer(idx1, method='bfill') assert_almost_equal(r1, rbfill1) assert_almost_equal(r2, rbfill2) # pass non-MultiIndex r1, r2 = idx1.get_indexer(idx2.get_tuple_index()) rexp1, rexp2 = idx1.get_indexer(idx2) assert_almost_equal(r1, rexp1) assert_almost_equal(r2, rexp2) self.assertRaises(Exception, idx1.get_indexer, list(zip(*idx2.get_tuple_index())[0])) def test_format(self): self.index.format() self.index[:0].format() def test_bounds(self): self.index._bounds def test_makeMask(self): from pandas.core.panel import make_mask mask = make_mask(self.index) expected = np.array([True, True, True, False, False, True, True, True], dtype=bool) self.assert_(np.array_equal(mask, expected)) def test_equals(self): self.assert_(self.index.equals(self.index)) self.assert_(self.index.equal_levels(self.index)) self.assert_(not self.index.equals(self.index[:-1])) self.assert_(not self.index.equals(self.index.get_tuple_index())) # different number of levels index = MultiIndex(levels=self.index.levels[:-1], labels=self.index.labels[:-1]) self.assert_(not self.index.equals(index)) self.assert_(not self.index.equal_levels(index)) # levels are different major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 2, 3]) minor_labels = np.array([0, 1, 0, 0, 1, 0]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assert_(not self.index.equals(index)) self.assert_(not self.index.equal_levels(index)) # some of the labels are different major_axis = Index(['foo', 'bar', 'baz', 'qux']) minor_axis = Index(['one', 'two']) major_labels = np.array([0, 0, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assert_(not self.index.equals(index)) def test_union(self): piece1 = self.index[:5][::-1] piece2 = self.index[3:] the_union = piece1.union(piece2) tups = sorted(self.index.get_tuple_index()) expected = MultiIndex.from_tuples(tups) self.assert_(the_union.equals(expected)) # corner case, pass self or empty thing: the_union = self.index.union(self.index) self.assert_(the_union is self.index) the_union = self.index.union(self.index[:0]) self.assert_(the_union is self.index) self.assertRaises(TypeError, self.index.union, self.index.get_tuple_index()) def test_intersection(self): piece1 = self.index[:5][::-1] piece2 = self.index[3:] the_int = piece1.intersection(piece2) tups = sorted(self.index[3:5].get_tuple_index()) expected = MultiIndex.from_tuples(tups) self.assert_(the_int.equals(expected)) # corner case, pass self the_int = self.index.intersection(self.index) self.assert_(the_int is self.index) self.assertRaises(TypeError, self.index.intersection, self.index.get_tuple_index()) def test_argsort(self): result = self.index.argsort() expected = self.index.get_tuple_index().argsort() self.assert_(np.array_equal(result, expected)) def test_sortlevel(self): import random tuples = list(self.index) random.shuffle(tuples) index = MultiIndex.from_tuples(tuples) sorted_idx, _ = index.sortlevel(0) expected = MultiIndex.from_tuples(sorted(tuples)) self.assert_(sorted_idx.equals(expected)) sorted_idx, _ = index.sortlevel(0, ascending=False) self.assert_(sorted_idx.equals(expected[::-1])) sorted_idx, _ = index.sortlevel(1) by1 = sorted(tuples, key=lambda x: (x[1], x[0])) expected = MultiIndex.from_tuples(by1) self.assert_(sorted_idx.equals(expected)) sorted_idx, _ = index.sortlevel(1, ascending=False) self.assert_(sorted_idx.equals(expected[::-1])) def test_dims(self): pass def test_drop(self): dropped = self.index.drop([('foo', 'two'), ('qux', 'one')]) index = MultiIndex.from_tuples([('foo', 'two'), ('qux', 'one')]) dropped2 = self.index.drop(index) expected = self.index[[0, 2, 3, 5]] self.assert_(dropped.equals(expected)) self.assert_(dropped2.equals(expected)) dropped = self.index.drop(['bar']) expected = self.index[[0, 1, 3, 4, 5]] self.assert_(dropped.equals(expected)) index = MultiIndex.from_tuples([('bar', 'two')]) self.assertRaises(Exception, self.index.drop, [('bar', 'two')]) self.assertRaises(Exception, self.index.drop, index) def test_insert(self): # key contained in all levels new_index = self.index.insert(0, ('bar', 'two')) self.assert_(new_index.equal_levels(self.index)) self.assert_(new_index[0] == ('bar', 'two')) # key not contained in all levels new_index = self.index.insert(0, ('abc', 'three')) self.assert_(np.array_equal(new_index.levels[0], list(self.index.levels[0]) + ['abc'])) self.assert_(np.array_equal(new_index.levels[1], list(self.index.levels[1]) + ['three'])) self.assert_(new_index[0] == ('abc', 'three')) # key wrong length self.assertRaises(Exception, self.index.insert, 0, ('foo2',))
class TestMultiIndex(unittest.TestCase): def setUp(self): major_axis = Index(['foo', 'bar', 'baz', 'qux']) minor_axis = Index(['one', 'two']) major_labels = np.array([0, 0, 1, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) self.index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels], names=['first', 'second']) def test_constructor_single_level(self): single_level = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux']], labels=[[0, 1, 2, 3]], names=['first']) self.assert_(isinstance(single_level, Index)) self.assert_(not isinstance(single_level, MultiIndex)) self.assert_(single_level.name == 'first') single_level = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux']], labels=[[0, 1, 2, 3]]) self.assert_(single_level.name is None) def test_constructor_no_levels(self): self.assertRaises(Exception, MultiIndex, levels=[], labels=[]) def test_duplicate_names(self): self.index.names = ['foo', 'foo'] self.assertRaises(Exception, self.index._get_level_number, 'foo') def test_from_arrays(self): arrays = [] for lev, lab in zip(self.index.levels, self.index.labels): arrays.append(np.asarray(lev).take(lab)) result = MultiIndex.from_arrays(arrays) self.assertEquals(list(result), list(self.index)) def test_append(self): result = self.index[:3].append(self.index[3:]) self.assert_(result.equals(self.index)) foos = [self.index[:1], self.index[1:3], self.index[3:]] result = foos[0].append(foos[1:]) self.assert_(result.equals(self.index)) # empty result = self.index.append([]) self.assert_(result.equals(self.index)) def test_get_level_values(self): result = self.index.get_level_values(0) expected = ['foo', 'foo', 'bar', 'baz', 'qux', 'qux'] self.assert_(np.array_equal(result, expected)) result = self.index.get_level_values('first') expected = self.index.get_level_values(0) self.assert_(np.array_equal(result, expected)) def test_nlevels(self): self.assertEquals(self.index.nlevels, 2) def test_iter(self): result = list(self.index) expected = [('foo', 'one'), ('foo', 'two'), ('bar', 'one'), ('baz', 'two'), ('qux', 'one'), ('qux', 'two')] self.assert_(result == expected) def test_pickle(self): pickled = pickle.dumps(self.index) unpickled = pickle.loads(pickled) self.assert_(self.index.equals(unpickled)) def test_legacy_pickle(self): import os def curpath(): pth, _ = os.path.split(os.path.abspath(__file__)) return pth ppath = os.path.join(curpath(), 'data/multiindex_v1.pickle') obj = pickle.load(open(ppath, 'r')) self.assert_(obj._is_legacy_format) obj2 = MultiIndex.from_tuples(obj.values) self.assert_(obj.equals(obj2)) res = obj.get_indexer(obj2[::-1]) exp = obj.get_indexer(obj[::-1]) exp2 = obj2.get_indexer(obj2[::-1]) assert_almost_equal(res, exp) assert_almost_equal(exp, exp2) def test_contains(self): self.assert_(('foo', 'two') in self.index) self.assert_(('bar', 'two') not in self.index) self.assert_(None not in self.index) def test_is_all_dates(self): self.assert_(not self.index.is_all_dates) def test_getitem(self): # scalar self.assertEquals(self.index[2], ('bar', 'one')) # slice result = self.index[2:5] expected = self.index[[2, 3, 4]] self.assert_(result.equals(expected)) # boolean result = self.index[[True, False, True, False, True, True]] result2 = self.index[np.array([True, False, True, False, True, True])] expected = self.index[[0, 2, 4, 5]] self.assert_(result.equals(expected)) self.assert_(result2.equals(expected)) def test_getitem_group_select(self): sorted_idx, _ = self.index.sortlevel(0) self.assertEquals(sorted_idx.get_loc('baz'), slice(3, 4)) self.assertEquals(sorted_idx.get_loc('foo'), slice(0, 2)) def test_get_loc(self): self.assert_(self.index.get_loc(('foo', 'two')) == 1) self.assert_(self.index.get_loc(('baz', 'two')) == 3) self.assertRaises(KeyError, self.index.get_loc, ('bar', 'two')) self.assertRaises(KeyError, self.index.get_loc, 'quux') # 3 levels index = MultiIndex( levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[ np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0]) ]) self.assertRaises(KeyError, index.get_loc, (1, 1)) self.assert_(index.get_loc((2, 0)) == slice(3, 5)) def test_get_loc_duplicates(self): index = Index([2, 2, 2, 2]) self.assertRaises(Exception, index.get_loc, 2) def test_slice_locs(self): df = tm.makeTimeDataFrame() stacked = df.stack() idx = stacked.index slob = slice(*idx.slice_locs(df.index[5], df.index[15])) sliced = stacked[slob] expected = df[5:16].stack() tm.assert_almost_equal(sliced.values, expected.values) slob = slice(*idx.slice_locs(df.index[5] + timedelta(seconds=30), df.index[15] - timedelta(seconds=30))) sliced = stacked[slob] expected = df[6:15].stack() tm.assert_almost_equal(sliced.values, expected.values) def test_slice_locs_not_sorted(self): index = MultiIndex( levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[ np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0]) ]) self.assertRaises(Exception, index.slice_locs, (1, 0, 1), (2, 1, 0)) # works sorted_index, _ = index.sortlevel(0) result = sorted_index.slice_locs((1, 0, 1), (2, 1, 0)) def test_slice_locs_partial(self): sorted_idx, _ = self.index.sortlevel(0) result = sorted_idx.slice_locs(('foo', 'two'), ('qux', 'one')) self.assertEquals(result, (1, 5)) result = sorted_idx.slice_locs(None, ('qux', 'one')) self.assertEquals(result, (0, 5)) result = sorted_idx.slice_locs(('foo', 'two'), None) self.assertEquals(result, (1, len(sorted_idx))) result = sorted_idx.slice_locs('bar', 'baz') self.assertEquals(result, (2, 4)) def test_slice_locs_not_contained(self): # some searchsorted action index = MultiIndex(levels=[[0, 2, 4, 6], [0, 2, 4]], labels=[[0, 0, 0, 1, 1, 2, 3, 3, 3], [0, 1, 2, 1, 2, 2, 0, 1, 2]], sortorder=0) result = index.slice_locs((1, 0), (5, 2)) self.assertEquals(result, (3, 6)) result = index.slice_locs(1, 5) self.assertEquals(result, (3, 6)) result = index.slice_locs((2, 2), (5, 2)) self.assertEquals(result, (3, 6)) result = index.slice_locs(2, 5) self.assertEquals(result, (3, 6)) result = index.slice_locs((1, 0), (6, 3)) self.assertEquals(result, (3, 8)) result = index.slice_locs(-1, 10) self.assertEquals(result, (0, len(index))) def test_consistency(self): # need to construct an overflow major_axis = range(70000) minor_axis = range(10) major_labels = np.arange(70000) minor_labels = np.repeat(range(10), 7000) # the fact that is works means it's consistent index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) # inconsistent major_labels = np.array([0, 0, 1, 1, 1, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assertRaises(Exception, getattr, index, 'indexMap') def test_truncate(self): major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) result = index.truncate(before=1) self.assert_('foo' not in result.levels[0]) self.assert_(1 in result.levels[0]) result = index.truncate(after=1) self.assert_(2 not in result.levels[0]) self.assert_(1 in result.levels[0]) result = index.truncate(before=1, after=2) self.assertEqual(len(result.levels[0]), 2) # after < before self.assertRaises(ValueError, index.truncate, 3, 1) def test_get_indexer(self): major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) idx1 = index[:5] idx2 = index[[1, 3, 5]] r1 = idx1.get_indexer(idx2) assert_almost_equal(r1, [1, 3, -1]) r1 = idx2.get_indexer(idx1, method='pad') assert_almost_equal(r1, [-1, 0, 0, 1, 1]) rffill1 = idx2.get_indexer(idx1, method='ffill') assert_almost_equal(r1, rffill1) r1 = idx2.get_indexer(idx1, method='backfill') assert_almost_equal(r1, [0, 0, 1, 1, 2]) rbfill1 = idx2.get_indexer(idx1, method='bfill') assert_almost_equal(r1, rbfill1) # pass non-MultiIndex r1 = idx1.get_indexer(idx2.get_tuple_index()) rexp1 = idx1.get_indexer(idx2) assert_almost_equal(r1, rexp1) # self.assertRaises(Exception, idx1.get_indexer, # list(list(zip(*idx2.get_tuple_index()))[0])) def test_format(self): self.index.format() self.index[:0].format() def test_format_integer_names(self): index = MultiIndex(levels=[[0, 1], [0, 1]], labels=[[0, 0, 1, 1], [0, 1, 0, 1]], names=[0, 1]) index.format(names=True) def test_bounds(self): self.index._bounds def test_equals(self): self.assert_(self.index.equals(self.index)) self.assert_(self.index.equal_levels(self.index)) self.assert_(not self.index.equals(self.index[:-1])) self.assert_(self.index.equals(self.index.get_tuple_index())) # different number of levels index = MultiIndex( levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[ np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0]) ]) index2 = MultiIndex(levels=index.levels[:-1], labels=index.labels[:-1]) self.assert_(not index.equals(index2)) self.assert_(not index.equal_levels(index2)) # levels are different major_axis = Index(range(4)) minor_axis = Index(range(2)) major_labels = np.array([0, 0, 1, 2, 2, 3]) minor_labels = np.array([0, 1, 0, 0, 1, 0]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assert_(not self.index.equals(index)) self.assert_(not self.index.equal_levels(index)) # some of the labels are different major_axis = Index(['foo', 'bar', 'baz', 'qux']) minor_axis = Index(['one', 'two']) major_labels = np.array([0, 0, 2, 2, 3, 3]) minor_labels = np.array([0, 1, 0, 1, 0, 1]) index = MultiIndex(levels=[major_axis, minor_axis], labels=[major_labels, minor_labels]) self.assert_(not self.index.equals(index)) def test_union(self): piece1 = self.index[:5][::-1] piece2 = self.index[3:] the_union = piece1 | piece2 tups = sorted(self.index.get_tuple_index()) expected = MultiIndex.from_tuples(tups) self.assert_(the_union.equals(expected)) # corner case, pass self or empty thing: the_union = self.index.union(self.index) self.assert_(the_union is self.index) the_union = self.index.union(self.index[:0]) self.assert_(the_union is self.index) # won't work in python 3 # tuples = self.index.get_tuple_index() # result = self.index[:4] | tuples[4:] # self.assert_(result.equals(tuples)) # not valid for python 3 # def test_union_with_regular_index(self): # other = Index(['A', 'B', 'C']) # result = other.union(self.index) # self.assert_(('foo', 'one') in result) # self.assert_('B' in result) # result2 = self.index.union(other) # self.assert_(result.equals(result2)) def test_intersection(self): piece1 = self.index[:5][::-1] piece2 = self.index[3:] the_int = piece1 & piece2 tups = sorted(self.index[3:5].get_tuple_index()) expected = MultiIndex.from_tuples(tups) self.assert_(the_int.equals(expected)) # corner case, pass self the_int = self.index.intersection(self.index) self.assert_(the_int is self.index) # empty intersection: disjoint empty = self.index[:2] & self.index[2:] expected = self.index[:0] self.assert_(empty.equals(expected)) # can't do in python 3 # tuples = self.index.get_tuple_index() # result = self.index & tuples # self.assert_(result.equals(tuples)) def test_diff(self): first = self.index result = first - self.index[-3:] expected = MultiIndex.from_tuples(sorted(self.index[:-3].values), sortorder=0, names=self.index.names) self.assert_(isinstance(result, MultiIndex)) self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # empty difference: reflexive result = self.index - self.index expected = self.index[:0] self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # empty difference: superset result = self.index[-3:] - self.index expected = self.index[:0] self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # empty difference: degenerate result = self.index[:0] - self.index expected = self.index[:0] self.assert_(result.equals(expected)) self.assertEqual(result.names, self.index.names) # names not the same chunklet = self.index[-3:] chunklet.names = ['foo', 'baz'] result = first - chunklet self.assertEqual(result.names, [None, None]) # empty, but non-equal result = self.index - self.index.sortlevel(1)[0] self.assert_(len(result) == 0) # raise Exception called with non-MultiIndex self.assertRaises(Exception, first.diff, first.get_tuple_index()) def test_from_tuples(self): self.assertRaises(Exception, MultiIndex.from_tuples, []) def test_argsort(self): result = self.index.argsort() expected = self.index.get_tuple_index().argsort() self.assert_(np.array_equal(result, expected)) def test_sortlevel(self): import random tuples = list(self.index) random.shuffle(tuples) index = MultiIndex.from_tuples(tuples) sorted_idx, _ = index.sortlevel(0) expected = MultiIndex.from_tuples(sorted(tuples)) self.assert_(sorted_idx.equals(expected)) sorted_idx, _ = index.sortlevel(0, ascending=False) self.assert_(sorted_idx.equals(expected[::-1])) sorted_idx, _ = index.sortlevel(1) by1 = sorted(tuples, key=lambda x: (x[1], x[0])) expected = MultiIndex.from_tuples(by1) self.assert_(sorted_idx.equals(expected)) sorted_idx, _ = index.sortlevel(1, ascending=False) self.assert_(sorted_idx.equals(expected[::-1])) def test_dims(self): pass def test_drop(self): dropped = self.index.drop([('foo', 'two'), ('qux', 'one')]) index = MultiIndex.from_tuples([('foo', 'two'), ('qux', 'one')]) dropped2 = self.index.drop(index) expected = self.index[[0, 2, 3, 5]] self.assert_(dropped.equals(expected)) self.assert_(dropped2.equals(expected)) dropped = self.index.drop(['bar']) expected = self.index[[0, 1, 3, 4, 5]] self.assert_(dropped.equals(expected)) index = MultiIndex.from_tuples([('bar', 'two')]) self.assertRaises(Exception, self.index.drop, [('bar', 'two')]) self.assertRaises(Exception, self.index.drop, index) # mixed partial / full drop dropped = self.index.drop(['foo', ('qux', 'one')]) expected = self.index[[2, 3, 5]] self.assert_(dropped.equals(expected)) def test_droplevel_with_names(self): index = self.index[self.index.get_loc('foo')] dropped = index.droplevel(0) self.assertEqual(dropped.name, 'second') index = MultiIndex( levels=[Index(range(4)), Index(range(4)), Index(range(4))], labels=[ np.array([0, 0, 1, 2, 2, 2, 3, 3]), np.array([0, 1, 0, 0, 0, 1, 0, 1]), np.array([1, 0, 1, 1, 0, 0, 1, 0]) ], names=['one', 'two', 'three']) dropped = index.droplevel(0) self.assertEqual(dropped.names, ['two', 'three']) def test_insert(self): # key contained in all levels new_index = self.index.insert(0, ('bar', 'two')) self.assert_(new_index.equal_levels(self.index)) self.assert_(new_index[0] == ('bar', 'two')) # key not contained in all levels new_index = self.index.insert(0, ('abc', 'three')) self.assert_( np.array_equal(new_index.levels[0], list(self.index.levels[0]) + ['abc'])) self.assert_( np.array_equal(new_index.levels[1], list(self.index.levels[1]) + ['three'])) self.assert_(new_index[0] == ('abc', 'three')) # key wrong length self.assertRaises(Exception, self.index.insert, 0, ('foo2', )) def test_take_preserve_name(self): taken = self.index.take([3, 0, 1]) self.assertEqual(taken.names, self.index.names) def test_join_level(self): def _check_how(other, how): join_index, lidx, ridx = other.join(self.index, how=how, level='second', return_indexers=True) exp_level = other.join(self.index.levels[1], how=how) self.assert_(join_index.levels[0].equals(self.index.levels[0])) self.assert_(join_index.levels[1].equals(exp_level)) # pare down levels mask = np.array([x[1] in exp_level for x in self.index], dtype=bool) exp_values = self.index.values[mask] self.assert_(np.array_equal(join_index.values, exp_values)) if how in ('outer', 'inner'): join_index2, ridx2, lidx2 = \ self.index.join(other, how=how, level='second', return_indexers=True) self.assert_(join_index.equals(join_index2)) self.assert_(np.array_equal(lidx, lidx2)) self.assert_(np.array_equal(ridx, ridx2)) self.assert_(np.array_equal(join_index2.values, exp_values)) def _check_all(other): _check_how(other, 'outer') _check_how(other, 'inner') _check_how(other, 'left') _check_how(other, 'right') _check_all(Index(['three', 'one', 'two'])) _check_all(Index(['one'])) _check_all(Index(['one', 'three'])) # some corner cases idx = Index(['three', 'one', 'two']) result = idx.join(self.index, level='second') self.assert_(isinstance(result, MultiIndex)) self.assertRaises(Exception, self.index.join, self.index, level=1) def test_reindex(self): result, indexer = self.index.reindex(list(self.index[:4])) self.assert_(isinstance(result, MultiIndex)) result, indexer = self.index.reindex(list(self.index)) self.assert_(isinstance(result, MultiIndex)) self.assert_(indexer is None) def test_reindex_level(self): idx = Index(['one']) target, indexer = self.index.reindex(idx, level='second') target2, indexer2 = idx.reindex(self.index, idx, level='second') exp_index = self.index.join(idx, level='second', how='left') self.assert_(target.equals(exp_index)) self.assert_(target2.equals(exp_index)) def test_has_duplicates(self): self.assert_(not self.index.has_duplicates) self.assert_(self.index.append(self.index).has_duplicates) index = MultiIndex(levels=[[0, 1], [0, 1, 2]], labels=[[0, 0, 0, 0, 1, 1, 1], [0, 1, 2, 0, 0, 1, 2]]) self.assert_(index.has_duplicates)