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)
self.assertRaises(Exception, self.index._get_level_number, "fourth")
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_reorder_levels(self):
# this blows up
self.assertRaises(Exception, self.index.reorder_levels, [2, 1, 0])
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_is_numeric(self):
# MultiIndex is never numeric
self.assert_(not self.index.is_numeric())
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, [])
idx = MultiIndex.from_tuples(((1, 2), (3, 4)), names=["a", "b"])
self.assertEquals(len(idx), 2)
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_sortlevel_deterministic(self):
tuples = [("bar", "one"), ("foo", "two"), ("qux", "two"), ("foo", "one"), ("baz", "two"), ("qux", "one")]
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, level="second")
exp_index = self.index.join(idx, level="second", how="right")
exp_index2 = self.index.join(idx, level="second", how="left")
self.assert_(target.equals(exp_index))
exp_indexer = np.array([0, 2, 4])
self.assert_(np.array_equal(indexer, exp_indexer))
self.assert_(target2.equals(exp_index2))
exp_indexer2 = np.array([0, -1, 0, -1, 0, -1])
self.assert_(np.array_equal(indexer2, exp_indexer2))
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)
