PageRenderTime 259ms CodeModel.GetById 27ms RepoModel.GetById 1ms app.codeStats 1ms

/pandas/tests/test_multilevel.py

http://github.com/wesm/pandas
Python | 2055 lines | 1908 code | 113 blank | 34 comment | 8 complexity | 652a9e533a6a899bdb5497d6302a19e9 MD5 | raw file
Possible License(s): BSD-3-Clause, Apache-2.0

Large files files are truncated, but you can click here to view the full file

  1. # -*- coding: utf-8 -*-
  2. # pylint: disable-msg=W0612,E1101,W0141
  3. import datetime
  4. import itertools
  5. from warnings import catch_warnings, simplefilter
  6. import numpy as np
  7. from numpy.random import randn
  8. import pytest
  9. import pytz
  10. from pandas.compat import (
  11. StringIO, lrange, lzip, product as cart_product, range, u, zip)
  12. from pandas.core.dtypes.common import is_float_dtype, is_integer_dtype
  13. import pandas as pd
  14. from pandas import DataFrame, Series, Timestamp, isna
  15. from pandas.core.index import Index, MultiIndex
  16. import pandas.util.testing as tm
  17. AGG_FUNCTIONS = ['sum', 'prod', 'min', 'max', 'median', 'mean', 'skew', 'mad',
  18. 'std', 'var', 'sem']
  19. class Base(object):
  20. def setup_method(self, method):
  21. index = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux'], ['one', 'two',
  22. 'three']],
  23. codes=[[0, 0, 0, 1, 1, 2, 2, 3, 3, 3],
  24. [0, 1, 2, 0, 1, 1, 2, 0, 1, 2]],
  25. names=['first', 'second'])
  26. self.frame = DataFrame(np.random.randn(10, 3), index=index,
  27. columns=Index(['A', 'B', 'C'], name='exp'))
  28. self.single_level = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux']],
  29. codes=[[0, 1, 2, 3]], names=['first'])
  30. # create test series object
  31. arrays = [['bar', 'bar', 'baz', 'baz', 'qux', 'qux', 'foo', 'foo'],
  32. ['one', 'two', 'one', 'two', 'one', 'two', 'one', 'two']]
  33. tuples = lzip(*arrays)
  34. index = MultiIndex.from_tuples(tuples)
  35. s = Series(randn(8), index=index)
  36. s[3] = np.NaN
  37. self.series = s
  38. self.tdf = tm.makeTimeDataFrame(100)
  39. self.ymd = self.tdf.groupby([lambda x: x.year, lambda x: x.month,
  40. lambda x: x.day]).sum()
  41. # use Int64Index, to make sure things work
  42. self.ymd.index.set_levels([lev.astype('i8')
  43. for lev in self.ymd.index.levels],
  44. inplace=True)
  45. self.ymd.index.set_names(['year', 'month', 'day'], inplace=True)
  46. class TestMultiLevel(Base):
  47. def test_append(self):
  48. a, b = self.frame[:5], self.frame[5:]
  49. result = a.append(b)
  50. tm.assert_frame_equal(result, self.frame)
  51. result = a['A'].append(b['A'])
  52. tm.assert_series_equal(result, self.frame['A'])
  53. def test_append_index(self):
  54. idx1 = Index([1.1, 1.2, 1.3])
  55. idx2 = pd.date_range('2011-01-01', freq='D', periods=3,
  56. tz='Asia/Tokyo')
  57. idx3 = Index(['A', 'B', 'C'])
  58. midx_lv2 = MultiIndex.from_arrays([idx1, idx2])
  59. midx_lv3 = MultiIndex.from_arrays([idx1, idx2, idx3])
  60. result = idx1.append(midx_lv2)
  61. # see gh-7112
  62. tz = pytz.timezone('Asia/Tokyo')
  63. expected_tuples = [(1.1, tz.localize(datetime.datetime(2011, 1, 1))),
  64. (1.2, tz.localize(datetime.datetime(2011, 1, 2))),
  65. (1.3, tz.localize(datetime.datetime(2011, 1, 3)))]
  66. expected = Index([1.1, 1.2, 1.3] + expected_tuples)
  67. tm.assert_index_equal(result, expected)
  68. result = midx_lv2.append(idx1)
  69. expected = Index(expected_tuples + [1.1, 1.2, 1.3])
  70. tm.assert_index_equal(result, expected)
  71. result = midx_lv2.append(midx_lv2)
  72. expected = MultiIndex.from_arrays([idx1.append(idx1),
  73. idx2.append(idx2)])
  74. tm.assert_index_equal(result, expected)
  75. result = midx_lv2.append(midx_lv3)
  76. tm.assert_index_equal(result, expected)
  77. result = midx_lv3.append(midx_lv2)
  78. expected = Index._simple_new(
  79. np.array([(1.1, tz.localize(datetime.datetime(2011, 1, 1)), 'A'),
  80. (1.2, tz.localize(datetime.datetime(2011, 1, 2)), 'B'),
  81. (1.3, tz.localize(datetime.datetime(2011, 1, 3)), 'C')] +
  82. expected_tuples), None)
  83. tm.assert_index_equal(result, expected)
  84. def test_dataframe_constructor(self):
  85. multi = DataFrame(np.random.randn(4, 4),
  86. index=[np.array(['a', 'a', 'b', 'b']),
  87. np.array(['x', 'y', 'x', 'y'])])
  88. assert isinstance(multi.index, MultiIndex)
  89. assert not isinstance(multi.columns, MultiIndex)
  90. multi = DataFrame(np.random.randn(4, 4),
  91. columns=[['a', 'a', 'b', 'b'],
  92. ['x', 'y', 'x', 'y']])
  93. assert isinstance(multi.columns, MultiIndex)
  94. def test_series_constructor(self):
  95. multi = Series(1., index=[np.array(['a', 'a', 'b', 'b']), np.array(
  96. ['x', 'y', 'x', 'y'])])
  97. assert isinstance(multi.index, MultiIndex)
  98. multi = Series(1., index=[['a', 'a', 'b', 'b'], ['x', 'y', 'x', 'y']])
  99. assert isinstance(multi.index, MultiIndex)
  100. multi = Series(lrange(4), index=[['a', 'a', 'b', 'b'],
  101. ['x', 'y', 'x', 'y']])
  102. assert isinstance(multi.index, MultiIndex)
  103. def test_reindex_level(self):
  104. # axis=0
  105. month_sums = self.ymd.sum(level='month')
  106. result = month_sums.reindex(self.ymd.index, level=1)
  107. expected = self.ymd.groupby(level='month').transform(np.sum)
  108. tm.assert_frame_equal(result, expected)
  109. # Series
  110. result = month_sums['A'].reindex(self.ymd.index, level=1)
  111. expected = self.ymd['A'].groupby(level='month').transform(np.sum)
  112. tm.assert_series_equal(result, expected, check_names=False)
  113. # axis=1
  114. month_sums = self.ymd.T.sum(axis=1, level='month')
  115. result = month_sums.reindex(columns=self.ymd.index, level=1)
  116. expected = self.ymd.groupby(level='month').transform(np.sum).T
  117. tm.assert_frame_equal(result, expected)
  118. def test_binops_level(self):
  119. def _check_op(opname):
  120. op = getattr(DataFrame, opname)
  121. month_sums = self.ymd.sum(level='month')
  122. result = op(self.ymd, month_sums, level='month')
  123. broadcasted = self.ymd.groupby(level='month').transform(np.sum)
  124. expected = op(self.ymd, broadcasted)
  125. tm.assert_frame_equal(result, expected)
  126. # Series
  127. op = getattr(Series, opname)
  128. result = op(self.ymd['A'], month_sums['A'], level='month')
  129. broadcasted = self.ymd['A'].groupby(level='month').transform(
  130. np.sum)
  131. expected = op(self.ymd['A'], broadcasted)
  132. expected.name = 'A'
  133. tm.assert_series_equal(result, expected)
  134. _check_op('sub')
  135. _check_op('add')
  136. _check_op('mul')
  137. _check_op('div')
  138. def test_pickle(self):
  139. def _test_roundtrip(frame):
  140. unpickled = tm.round_trip_pickle(frame)
  141. tm.assert_frame_equal(frame, unpickled)
  142. _test_roundtrip(self.frame)
  143. _test_roundtrip(self.frame.T)
  144. _test_roundtrip(self.ymd)
  145. _test_roundtrip(self.ymd.T)
  146. def test_reindex(self):
  147. expected = self.frame.iloc[[0, 3]]
  148. reindexed = self.frame.loc[[('foo', 'one'), ('bar', 'one')]]
  149. tm.assert_frame_equal(reindexed, expected)
  150. with catch_warnings(record=True):
  151. simplefilter("ignore", DeprecationWarning)
  152. reindexed = self.frame.ix[[('foo', 'one'), ('bar', 'one')]]
  153. tm.assert_frame_equal(reindexed, expected)
  154. def test_reindex_preserve_levels(self):
  155. new_index = self.ymd.index[::10]
  156. chunk = self.ymd.reindex(new_index)
  157. assert chunk.index is new_index
  158. chunk = self.ymd.loc[new_index]
  159. assert chunk.index is new_index
  160. with catch_warnings(record=True):
  161. simplefilter("ignore", DeprecationWarning)
  162. chunk = self.ymd.ix[new_index]
  163. assert chunk.index is new_index
  164. ymdT = self.ymd.T
  165. chunk = ymdT.reindex(columns=new_index)
  166. assert chunk.columns is new_index
  167. chunk = ymdT.loc[:, new_index]
  168. assert chunk.columns is new_index
  169. def test_repr_to_string(self):
  170. repr(self.frame)
  171. repr(self.ymd)
  172. repr(self.frame.T)
  173. repr(self.ymd.T)
  174. buf = StringIO()
  175. self.frame.to_string(buf=buf)
  176. self.ymd.to_string(buf=buf)
  177. self.frame.T.to_string(buf=buf)
  178. self.ymd.T.to_string(buf=buf)
  179. def test_repr_name_coincide(self):
  180. index = MultiIndex.from_tuples([('a', 0, 'foo'), ('b', 1, 'bar')],
  181. names=['a', 'b', 'c'])
  182. df = DataFrame({'value': [0, 1]}, index=index)
  183. lines = repr(df).split('\n')
  184. assert lines[2].startswith('a 0 foo')
  185. def test_delevel_infer_dtype(self):
  186. tuples = [tuple
  187. for tuple in cart_product(
  188. ['foo', 'bar'], [10, 20], [1.0, 1.1])]
  189. index = MultiIndex.from_tuples(tuples, names=['prm0', 'prm1', 'prm2'])
  190. df = DataFrame(np.random.randn(8, 3), columns=['A', 'B', 'C'],
  191. index=index)
  192. deleveled = df.reset_index()
  193. assert is_integer_dtype(deleveled['prm1'])
  194. assert is_float_dtype(deleveled['prm2'])
  195. def test_reset_index_with_drop(self):
  196. deleveled = self.ymd.reset_index(drop=True)
  197. assert len(deleveled.columns) == len(self.ymd.columns)
  198. assert deleveled.index.name == self.ymd.index.name
  199. deleveled = self.series.reset_index()
  200. assert isinstance(deleveled, DataFrame)
  201. assert len(deleveled.columns) == len(self.series.index.levels) + 1
  202. assert deleveled.index.name == self.series.index.name
  203. deleveled = self.series.reset_index(drop=True)
  204. assert isinstance(deleveled, Series)
  205. assert deleveled.index.name == self.series.index.name
  206. def test_count_level(self):
  207. def _check_counts(frame, axis=0):
  208. index = frame._get_axis(axis)
  209. for i in range(index.nlevels):
  210. result = frame.count(axis=axis, level=i)
  211. expected = frame.groupby(axis=axis, level=i).count()
  212. expected = expected.reindex_like(result).astype('i8')
  213. tm.assert_frame_equal(result, expected)
  214. self.frame.iloc[1, [1, 2]] = np.nan
  215. self.frame.iloc[7, [0, 1]] = np.nan
  216. self.ymd.iloc[1, [1, 2]] = np.nan
  217. self.ymd.iloc[7, [0, 1]] = np.nan
  218. _check_counts(self.frame)
  219. _check_counts(self.ymd)
  220. _check_counts(self.frame.T, axis=1)
  221. _check_counts(self.ymd.T, axis=1)
  222. # can't call with level on regular DataFrame
  223. df = tm.makeTimeDataFrame()
  224. with pytest.raises(TypeError, match='hierarchical'):
  225. df.count(level=0)
  226. self.frame['D'] = 'foo'
  227. result = self.frame.count(level=0, numeric_only=True)
  228. tm.assert_index_equal(result.columns, Index(list('ABC'), name='exp'))
  229. def test_count_level_series(self):
  230. index = MultiIndex(levels=[['foo', 'bar', 'baz'], ['one', 'two',
  231. 'three', 'four']],
  232. codes=[[0, 0, 0, 2, 2], [2, 0, 1, 1, 2]])
  233. s = Series(np.random.randn(len(index)), index=index)
  234. result = s.count(level=0)
  235. expected = s.groupby(level=0).count()
  236. tm.assert_series_equal(
  237. result.astype('f8'), expected.reindex(result.index).fillna(0))
  238. result = s.count(level=1)
  239. expected = s.groupby(level=1).count()
  240. tm.assert_series_equal(
  241. result.astype('f8'), expected.reindex(result.index).fillna(0))
  242. def test_count_level_corner(self):
  243. s = self.frame['A'][:0]
  244. result = s.count(level=0)
  245. expected = Series(0, index=s.index.levels[0], name='A')
  246. tm.assert_series_equal(result, expected)
  247. df = self.frame[:0]
  248. result = df.count(level=0)
  249. expected = DataFrame({}, index=s.index.levels[0],
  250. columns=df.columns).fillna(0).astype(np.int64)
  251. tm.assert_frame_equal(result, expected)
  252. def test_get_level_number_out_of_bounds(self):
  253. with pytest.raises(IndexError, match="Too many levels"):
  254. self.frame.index._get_level_number(2)
  255. with pytest.raises(IndexError, match="not a valid level number"):
  256. self.frame.index._get_level_number(-3)
  257. def test_unstack(self):
  258. # just check that it works for now
  259. unstacked = self.ymd.unstack()
  260. unstacked.unstack()
  261. # test that ints work
  262. self.ymd.astype(int).unstack()
  263. # test that int32 work
  264. self.ymd.astype(np.int32).unstack()
  265. def test_unstack_multiple_no_empty_columns(self):
  266. index = MultiIndex.from_tuples([(0, 'foo', 0), (0, 'bar', 0), (
  267. 1, 'baz', 1), (1, 'qux', 1)])
  268. s = Series(np.random.randn(4), index=index)
  269. unstacked = s.unstack([1, 2])
  270. expected = unstacked.dropna(axis=1, how='all')
  271. tm.assert_frame_equal(unstacked, expected)
  272. def test_stack(self):
  273. # regular roundtrip
  274. unstacked = self.ymd.unstack()
  275. restacked = unstacked.stack()
  276. tm.assert_frame_equal(restacked, self.ymd)
  277. unlexsorted = self.ymd.sort_index(level=2)
  278. unstacked = unlexsorted.unstack(2)
  279. restacked = unstacked.stack()
  280. tm.assert_frame_equal(restacked.sort_index(level=0), self.ymd)
  281. unlexsorted = unlexsorted[::-1]
  282. unstacked = unlexsorted.unstack(1)
  283. restacked = unstacked.stack().swaplevel(1, 2)
  284. tm.assert_frame_equal(restacked.sort_index(level=0), self.ymd)
  285. unlexsorted = unlexsorted.swaplevel(0, 1)
  286. unstacked = unlexsorted.unstack(0).swaplevel(0, 1, axis=1)
  287. restacked = unstacked.stack(0).swaplevel(1, 2)
  288. tm.assert_frame_equal(restacked.sort_index(level=0), self.ymd)
  289. # columns unsorted
  290. unstacked = self.ymd.unstack()
  291. unstacked = unstacked.sort_index(axis=1, ascending=False)
  292. restacked = unstacked.stack()
  293. tm.assert_frame_equal(restacked, self.ymd)
  294. # more than 2 levels in the columns
  295. unstacked = self.ymd.unstack(1).unstack(1)
  296. result = unstacked.stack(1)
  297. expected = self.ymd.unstack()
  298. tm.assert_frame_equal(result, expected)
  299. result = unstacked.stack(2)
  300. expected = self.ymd.unstack(1)
  301. tm.assert_frame_equal(result, expected)
  302. result = unstacked.stack(0)
  303. expected = self.ymd.stack().unstack(1).unstack(1)
  304. tm.assert_frame_equal(result, expected)
  305. # not all levels present in each echelon
  306. unstacked = self.ymd.unstack(2).loc[:, ::3]
  307. stacked = unstacked.stack().stack()
  308. ymd_stacked = self.ymd.stack()
  309. tm.assert_series_equal(stacked, ymd_stacked.reindex(stacked.index))
  310. # stack with negative number
  311. result = self.ymd.unstack(0).stack(-2)
  312. expected = self.ymd.unstack(0).stack(0)
  313. # GH10417
  314. def check(left, right):
  315. tm.assert_series_equal(left, right)
  316. assert left.index.is_unique is False
  317. li, ri = left.index, right.index
  318. tm.assert_index_equal(li, ri)
  319. df = DataFrame(np.arange(12).reshape(4, 3),
  320. index=list('abab'),
  321. columns=['1st', '2nd', '3rd'])
  322. mi = MultiIndex(levels=[['a', 'b'], ['1st', '2nd', '3rd']],
  323. codes=[np.tile(
  324. np.arange(2).repeat(3), 2), np.tile(
  325. np.arange(3), 4)])
  326. left, right = df.stack(), Series(np.arange(12), index=mi)
  327. check(left, right)
  328. df.columns = ['1st', '2nd', '1st']
  329. mi = MultiIndex(levels=[['a', 'b'], ['1st', '2nd']], codes=[np.tile(
  330. np.arange(2).repeat(3), 2), np.tile(
  331. [0, 1, 0], 4)])
  332. left, right = df.stack(), Series(np.arange(12), index=mi)
  333. check(left, right)
  334. tpls = ('a', 2), ('b', 1), ('a', 1), ('b', 2)
  335. df.index = MultiIndex.from_tuples(tpls)
  336. mi = MultiIndex(levels=[['a', 'b'], [1, 2], ['1st', '2nd']],
  337. codes=[np.tile(
  338. np.arange(2).repeat(3), 2), np.repeat(
  339. [1, 0, 1], [3, 6, 3]), np.tile(
  340. [0, 1, 0], 4)])
  341. left, right = df.stack(), Series(np.arange(12), index=mi)
  342. check(left, right)
  343. def test_unstack_odd_failure(self):
  344. data = """day,time,smoker,sum,len
  345. Fri,Dinner,No,8.25,3.
  346. Fri,Dinner,Yes,27.03,9
  347. Fri,Lunch,No,3.0,1
  348. Fri,Lunch,Yes,13.68,6
  349. Sat,Dinner,No,139.63,45
  350. Sat,Dinner,Yes,120.77,42
  351. Sun,Dinner,No,180.57,57
  352. Sun,Dinner,Yes,66.82,19
  353. Thur,Dinner,No,3.0,1
  354. Thur,Lunch,No,117.32,44
  355. Thur,Lunch,Yes,51.51,17"""
  356. df = pd.read_csv(StringIO(data)).set_index(['day', 'time', 'smoker'])
  357. # it works, #2100
  358. result = df.unstack(2)
  359. recons = result.stack()
  360. tm.assert_frame_equal(recons, df)
  361. def test_stack_mixed_dtype(self):
  362. df = self.frame.T
  363. df['foo', 'four'] = 'foo'
  364. df = df.sort_index(level=1, axis=1)
  365. stacked = df.stack()
  366. result = df['foo'].stack().sort_index()
  367. tm.assert_series_equal(stacked['foo'], result, check_names=False)
  368. assert result.name is None
  369. assert stacked['bar'].dtype == np.float_
  370. def test_unstack_bug(self):
  371. df = DataFrame({'state': ['naive', 'naive', 'naive', 'activ', 'activ',
  372. 'activ'],
  373. 'exp': ['a', 'b', 'b', 'b', 'a', 'a'],
  374. 'barcode': [1, 2, 3, 4, 1, 3],
  375. 'v': ['hi', 'hi', 'bye', 'bye', 'bye', 'peace'],
  376. 'extra': np.arange(6.)})
  377. result = df.groupby(['state', 'exp', 'barcode', 'v']).apply(len)
  378. unstacked = result.unstack()
  379. restacked = unstacked.stack()
  380. tm.assert_series_equal(
  381. restacked, result.reindex(restacked.index).astype(float))
  382. def test_stack_unstack_preserve_names(self):
  383. unstacked = self.frame.unstack()
  384. assert unstacked.index.name == 'first'
  385. assert unstacked.columns.names == ['exp', 'second']
  386. restacked = unstacked.stack()
  387. assert restacked.index.names == self.frame.index.names
  388. def test_unstack_level_name(self):
  389. result = self.frame.unstack('second')
  390. expected = self.frame.unstack(level=1)
  391. tm.assert_frame_equal(result, expected)
  392. def test_stack_level_name(self):
  393. unstacked = self.frame.unstack('second')
  394. result = unstacked.stack('exp')
  395. expected = self.frame.unstack().stack(0)
  396. tm.assert_frame_equal(result, expected)
  397. result = self.frame.stack('exp')
  398. expected = self.frame.stack()
  399. tm.assert_series_equal(result, expected)
  400. def test_stack_unstack_multiple(self):
  401. unstacked = self.ymd.unstack(['year', 'month'])
  402. expected = self.ymd.unstack('year').unstack('month')
  403. tm.assert_frame_equal(unstacked, expected)
  404. assert unstacked.columns.names == expected.columns.names
  405. # series
  406. s = self.ymd['A']
  407. s_unstacked = s.unstack(['year', 'month'])
  408. tm.assert_frame_equal(s_unstacked, expected['A'])
  409. restacked = unstacked.stack(['year', 'month'])
  410. restacked = restacked.swaplevel(0, 1).swaplevel(1, 2)
  411. restacked = restacked.sort_index(level=0)
  412. tm.assert_frame_equal(restacked, self.ymd)
  413. assert restacked.index.names == self.ymd.index.names
  414. # GH #451
  415. unstacked = self.ymd.unstack([1, 2])
  416. expected = self.ymd.unstack(1).unstack(1).dropna(axis=1, how='all')
  417. tm.assert_frame_equal(unstacked, expected)
  418. unstacked = self.ymd.unstack([2, 1])
  419. expected = self.ymd.unstack(2).unstack(1).dropna(axis=1, how='all')
  420. tm.assert_frame_equal(unstacked, expected.loc[:, unstacked.columns])
  421. def test_stack_names_and_numbers(self):
  422. unstacked = self.ymd.unstack(['year', 'month'])
  423. # Can't use mixture of names and numbers to stack
  424. with pytest.raises(ValueError, match="level should contain"):
  425. unstacked.stack([0, 'month'])
  426. def test_stack_multiple_out_of_bounds(self):
  427. # nlevels == 3
  428. unstacked = self.ymd.unstack(['year', 'month'])
  429. with pytest.raises(IndexError, match="Too many levels"):
  430. unstacked.stack([2, 3])
  431. with pytest.raises(IndexError, match="not a valid level number"):
  432. unstacked.stack([-4, -3])
  433. def test_unstack_period_series(self):
  434. # GH 4342
  435. idx1 = pd.PeriodIndex(['2013-01', '2013-01', '2013-02', '2013-02',
  436. '2013-03', '2013-03'], freq='M', name='period')
  437. idx2 = Index(['A', 'B'] * 3, name='str')
  438. value = [1, 2, 3, 4, 5, 6]
  439. idx = MultiIndex.from_arrays([idx1, idx2])
  440. s = Series(value, index=idx)
  441. result1 = s.unstack()
  442. result2 = s.unstack(level=1)
  443. result3 = s.unstack(level=0)
  444. e_idx = pd.PeriodIndex(
  445. ['2013-01', '2013-02', '2013-03'], freq='M', name='period')
  446. expected = DataFrame({'A': [1, 3, 5], 'B': [2, 4, 6]}, index=e_idx,
  447. columns=['A', 'B'])
  448. expected.columns.name = 'str'
  449. tm.assert_frame_equal(result1, expected)
  450. tm.assert_frame_equal(result2, expected)
  451. tm.assert_frame_equal(result3, expected.T)
  452. idx1 = pd.PeriodIndex(['2013-01', '2013-01', '2013-02', '2013-02',
  453. '2013-03', '2013-03'], freq='M', name='period1')
  454. idx2 = pd.PeriodIndex(['2013-12', '2013-11', '2013-10', '2013-09',
  455. '2013-08', '2013-07'], freq='M', name='period2')
  456. idx = MultiIndex.from_arrays([idx1, idx2])
  457. s = Series(value, index=idx)
  458. result1 = s.unstack()
  459. result2 = s.unstack(level=1)
  460. result3 = s.unstack(level=0)
  461. e_idx = pd.PeriodIndex(
  462. ['2013-01', '2013-02', '2013-03'], freq='M', name='period1')
  463. e_cols = pd.PeriodIndex(['2013-07', '2013-08', '2013-09', '2013-10',
  464. '2013-11', '2013-12'],
  465. freq='M', name='period2')
  466. expected = DataFrame([[np.nan, np.nan, np.nan, np.nan, 2, 1],
  467. [np.nan, np.nan, 4, 3, np.nan, np.nan],
  468. [6, 5, np.nan, np.nan, np.nan, np.nan]],
  469. index=e_idx, columns=e_cols)
  470. tm.assert_frame_equal(result1, expected)
  471. tm.assert_frame_equal(result2, expected)
  472. tm.assert_frame_equal(result3, expected.T)
  473. def test_unstack_period_frame(self):
  474. # GH 4342
  475. idx1 = pd.PeriodIndex(['2014-01', '2014-02', '2014-02', '2014-02',
  476. '2014-01', '2014-01'],
  477. freq='M', name='period1')
  478. idx2 = pd.PeriodIndex(['2013-12', '2013-12', '2014-02', '2013-10',
  479. '2013-10', '2014-02'],
  480. freq='M', name='period2')
  481. value = {'A': [1, 2, 3, 4, 5, 6], 'B': [6, 5, 4, 3, 2, 1]}
  482. idx = MultiIndex.from_arrays([idx1, idx2])
  483. df = DataFrame(value, index=idx)
  484. result1 = df.unstack()
  485. result2 = df.unstack(level=1)
  486. result3 = df.unstack(level=0)
  487. e_1 = pd.PeriodIndex(['2014-01', '2014-02'], freq='M', name='period1')
  488. e_2 = pd.PeriodIndex(['2013-10', '2013-12', '2014-02', '2013-10',
  489. '2013-12', '2014-02'], freq='M', name='period2')
  490. e_cols = MultiIndex.from_arrays(['A A A B B B'.split(), e_2])
  491. expected = DataFrame([[5, 1, 6, 2, 6, 1], [4, 2, 3, 3, 5, 4]],
  492. index=e_1, columns=e_cols)
  493. tm.assert_frame_equal(result1, expected)
  494. tm.assert_frame_equal(result2, expected)
  495. e_1 = pd.PeriodIndex(['2014-01', '2014-02', '2014-01',
  496. '2014-02'], freq='M', name='period1')
  497. e_2 = pd.PeriodIndex(
  498. ['2013-10', '2013-12', '2014-02'], freq='M', name='period2')
  499. e_cols = MultiIndex.from_arrays(['A A B B'.split(), e_1])
  500. expected = DataFrame([[5, 4, 2, 3], [1, 2, 6, 5], [6, 3, 1, 4]],
  501. index=e_2, columns=e_cols)
  502. tm.assert_frame_equal(result3, expected)
  503. def test_stack_multiple_bug(self):
  504. """ bug when some uniques are not present in the data #3170"""
  505. id_col = ([1] * 3) + ([2] * 3)
  506. name = (['a'] * 3) + (['b'] * 3)
  507. date = pd.to_datetime(['2013-01-03', '2013-01-04', '2013-01-05'] * 2)
  508. var1 = np.random.randint(0, 100, 6)
  509. df = DataFrame(dict(ID=id_col, NAME=name, DATE=date, VAR1=var1))
  510. multi = df.set_index(['DATE', 'ID'])
  511. multi.columns.name = 'Params'
  512. unst = multi.unstack('ID')
  513. down = unst.resample('W-THU').mean()
  514. rs = down.stack('ID')
  515. xp = unst.loc[:, ['VAR1']].resample('W-THU').mean().stack('ID')
  516. xp.columns.name = 'Params'
  517. tm.assert_frame_equal(rs, xp)
  518. def test_stack_dropna(self):
  519. # GH #3997
  520. df = DataFrame({'A': ['a1', 'a2'], 'B': ['b1', 'b2'], 'C': [1, 1]})
  521. df = df.set_index(['A', 'B'])
  522. stacked = df.unstack().stack(dropna=False)
  523. assert len(stacked) > len(stacked.dropna())
  524. stacked = df.unstack().stack(dropna=True)
  525. tm.assert_frame_equal(stacked, stacked.dropna())
  526. def test_unstack_multiple_hierarchical(self):
  527. df = DataFrame(index=[[0, 0, 0, 0, 1, 1, 1, 1],
  528. [0, 0, 1, 1, 0, 0, 1, 1], [0, 1, 0, 1, 0, 1, 0, 1
  529. ]],
  530. columns=[[0, 0, 1, 1], [0, 1, 0, 1]])
  531. df.index.names = ['a', 'b', 'c']
  532. df.columns.names = ['d', 'e']
  533. # it works!
  534. df.unstack(['b', 'c'])
  535. def test_groupby_transform(self):
  536. s = self.frame['A']
  537. grouper = s.index.get_level_values(0)
  538. grouped = s.groupby(grouper)
  539. applied = grouped.apply(lambda x: x * 2)
  540. expected = grouped.transform(lambda x: x * 2)
  541. result = applied.reindex(expected.index)
  542. tm.assert_series_equal(result, expected, check_names=False)
  543. def test_unstack_sparse_keyspace(self):
  544. # memory problems with naive impl #2278
  545. # Generate Long File & Test Pivot
  546. NUM_ROWS = 1000
  547. df = DataFrame({'A': np.random.randint(100, size=NUM_ROWS),
  548. 'B': np.random.randint(300, size=NUM_ROWS),
  549. 'C': np.random.randint(-7, 7, size=NUM_ROWS),
  550. 'D': np.random.randint(-19, 19, size=NUM_ROWS),
  551. 'E': np.random.randint(3000, size=NUM_ROWS),
  552. 'F': np.random.randn(NUM_ROWS)})
  553. idf = df.set_index(['A', 'B', 'C', 'D', 'E'])
  554. # it works! is sufficient
  555. idf.unstack('E')
  556. def test_unstack_unobserved_keys(self):
  557. # related to #2278 refactoring
  558. levels = [[0, 1], [0, 1, 2, 3]]
  559. codes = [[0, 0, 1, 1], [0, 2, 0, 2]]
  560. index = MultiIndex(levels, codes)
  561. df = DataFrame(np.random.randn(4, 2), index=index)
  562. result = df.unstack()
  563. assert len(result.columns) == 4
  564. recons = result.stack()
  565. tm.assert_frame_equal(recons, df)
  566. @pytest.mark.slow
  567. def test_unstack_number_of_levels_larger_than_int32(self):
  568. # GH 20601
  569. df = DataFrame(np.random.randn(2 ** 16, 2),
  570. index=[np.arange(2 ** 16), np.arange(2 ** 16)])
  571. with pytest.raises(ValueError, match='int32 overflow'):
  572. df.unstack()
  573. def test_stack_order_with_unsorted_levels(self):
  574. # GH 16323
  575. def manual_compare_stacked(df, df_stacked, lev0, lev1):
  576. assert all(df.loc[row, col] ==
  577. df_stacked.loc[(row, col[lev0]), col[lev1]]
  578. for row in df.index for col in df.columns)
  579. # deep check for 1-row case
  580. for width in [2, 3]:
  581. levels_poss = itertools.product(
  582. itertools.permutations([0, 1, 2], width),
  583. repeat=2)
  584. for levels in levels_poss:
  585. columns = MultiIndex(levels=levels,
  586. codes=[[0, 0, 1, 1],
  587. [0, 1, 0, 1]])
  588. df = DataFrame(columns=columns, data=[range(4)])
  589. for stack_lev in range(2):
  590. df_stacked = df.stack(stack_lev)
  591. manual_compare_stacked(df, df_stacked,
  592. stack_lev, 1 - stack_lev)
  593. # check multi-row case
  594. mi = MultiIndex(levels=[["A", "C", "B"], ["B", "A", "C"]],
  595. codes=[np.repeat(range(3), 3), np.tile(range(3), 3)])
  596. df = DataFrame(columns=mi, index=range(5),
  597. data=np.arange(5 * len(mi)).reshape(5, -1))
  598. manual_compare_stacked(df, df.stack(0), 0, 1)
  599. def test_groupby_corner(self):
  600. midx = MultiIndex(levels=[['foo'], ['bar'], ['baz']],
  601. codes=[[0], [0], [0]],
  602. names=['one', 'two', 'three'])
  603. df = DataFrame([np.random.rand(4)], columns=['a', 'b', 'c', 'd'],
  604. index=midx)
  605. # should work
  606. df.groupby(level='three')
  607. def test_groupby_level_no_obs(self):
  608. # #1697
  609. midx = MultiIndex.from_tuples([('f1', 's1'), ('f1', 's2'), (
  610. 'f2', 's1'), ('f2', 's2'), ('f3', 's1'), ('f3', 's2')])
  611. df = DataFrame(
  612. [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]], columns=midx)
  613. df1 = df.loc(axis=1)[df.columns.map(
  614. lambda u: u[0] in ['f2', 'f3'])]
  615. grouped = df1.groupby(axis=1, level=0)
  616. result = grouped.sum()
  617. assert (result.columns == ['f2', 'f3']).all()
  618. def test_join(self):
  619. a = self.frame.loc[self.frame.index[:5], ['A']]
  620. b = self.frame.loc[self.frame.index[2:], ['B', 'C']]
  621. joined = a.join(b, how='outer').reindex(self.frame.index)
  622. expected = self.frame.copy()
  623. expected.values[np.isnan(joined.values)] = np.nan
  624. assert not np.isnan(joined.values).all()
  625. # TODO what should join do with names ?
  626. tm.assert_frame_equal(joined, expected, check_names=False)
  627. def test_swaplevel(self):
  628. swapped = self.frame['A'].swaplevel()
  629. swapped2 = self.frame['A'].swaplevel(0)
  630. swapped3 = self.frame['A'].swaplevel(0, 1)
  631. swapped4 = self.frame['A'].swaplevel('first', 'second')
  632. assert not swapped.index.equals(self.frame.index)
  633. tm.assert_series_equal(swapped, swapped2)
  634. tm.assert_series_equal(swapped, swapped3)
  635. tm.assert_series_equal(swapped, swapped4)
  636. back = swapped.swaplevel()
  637. back2 = swapped.swaplevel(0)
  638. back3 = swapped.swaplevel(0, 1)
  639. back4 = swapped.swaplevel('second', 'first')
  640. assert back.index.equals(self.frame.index)
  641. tm.assert_series_equal(back, back2)
  642. tm.assert_series_equal(back, back3)
  643. tm.assert_series_equal(back, back4)
  644. ft = self.frame.T
  645. swapped = ft.swaplevel('first', 'second', axis=1)
  646. exp = self.frame.swaplevel('first', 'second').T
  647. tm.assert_frame_equal(swapped, exp)
  648. def test_reorder_levels(self):
  649. result = self.ymd.reorder_levels(['month', 'day', 'year'])
  650. expected = self.ymd.swaplevel(0, 1).swaplevel(1, 2)
  651. tm.assert_frame_equal(result, expected)
  652. result = self.ymd['A'].reorder_levels(['month', 'day', 'year'])
  653. expected = self.ymd['A'].swaplevel(0, 1).swaplevel(1, 2)
  654. tm.assert_series_equal(result, expected)
  655. result = self.ymd.T.reorder_levels(['month', 'day', 'year'], axis=1)
  656. expected = self.ymd.T.swaplevel(0, 1, axis=1).swaplevel(1, 2, axis=1)
  657. tm.assert_frame_equal(result, expected)
  658. with pytest.raises(TypeError, match='hierarchical axis'):
  659. self.ymd.reorder_levels([1, 2], axis=1)
  660. with pytest.raises(IndexError, match='Too many levels'):
  661. self.ymd.index.reorder_levels([1, 2, 3])
  662. def test_insert_index(self):
  663. df = self.ymd[:5].T
  664. df[2000, 1, 10] = df[2000, 1, 7]
  665. assert isinstance(df.columns, MultiIndex)
  666. assert (df[2000, 1, 10] == df[2000, 1, 7]).all()
  667. def test_alignment(self):
  668. x = Series(data=[1, 2, 3], index=MultiIndex.from_tuples([("A", 1), (
  669. "A", 2), ("B", 3)]))
  670. y = Series(data=[4, 5, 6], index=MultiIndex.from_tuples([("Z", 1), (
  671. "Z", 2), ("B", 3)]))
  672. res = x - y
  673. exp_index = x.index.union(y.index)
  674. exp = x.reindex(exp_index) - y.reindex(exp_index)
  675. tm.assert_series_equal(res, exp)
  676. # hit non-monotonic code path
  677. res = x[::-1] - y[::-1]
  678. exp_index = x.index.union(y.index)
  679. exp = x.reindex(exp_index) - y.reindex(exp_index)
  680. tm.assert_series_equal(res, exp)
  681. def test_count(self):
  682. frame = self.frame.copy()
  683. frame.index.names = ['a', 'b']
  684. result = frame.count(level='b')
  685. expect = self.frame.count(level=1)
  686. tm.assert_frame_equal(result, expect, check_names=False)
  687. result = frame.count(level='a')
  688. expect = self.frame.count(level=0)
  689. tm.assert_frame_equal(result, expect, check_names=False)
  690. series = self.series.copy()
  691. series.index.names = ['a', 'b']
  692. result = series.count(level='b')
  693. expect = self.series.count(level=1)
  694. tm.assert_series_equal(result, expect, check_names=False)
  695. assert result.index.name == 'b'
  696. result = series.count(level='a')
  697. expect = self.series.count(level=0)
  698. tm.assert_series_equal(result, expect, check_names=False)
  699. assert result.index.name == 'a'
  700. msg = "Level x not found"
  701. with pytest.raises(KeyError, match=msg):
  702. series.count('x')
  703. with pytest.raises(KeyError, match=msg):
  704. frame.count(level='x')
  705. @pytest.mark.parametrize('op', AGG_FUNCTIONS)
  706. @pytest.mark.parametrize('level', [0, 1])
  707. @pytest.mark.parametrize('skipna', [True, False])
  708. @pytest.mark.parametrize('sort', [True, False])
  709. def test_series_group_min_max(self, op, level, skipna, sort):
  710. # GH 17537
  711. grouped = self.series.groupby(level=level, sort=sort)
  712. # skipna=True
  713. leftside = grouped.agg(lambda x: getattr(x, op)(skipna=skipna))
  714. rightside = getattr(self.series, op)(level=level, skipna=skipna)
  715. if sort:
  716. rightside = rightside.sort_index(level=level)
  717. tm.assert_series_equal(leftside, rightside)
  718. @pytest.mark.parametrize('op', AGG_FUNCTIONS)
  719. @pytest.mark.parametrize('level', [0, 1])
  720. @pytest.mark.parametrize('axis', [0, 1])
  721. @pytest.mark.parametrize('skipna', [True, False])
  722. @pytest.mark.parametrize('sort', [True, False])
  723. def test_frame_group_ops(self, op, level, axis, skipna, sort):
  724. # GH 17537
  725. self.frame.iloc[1, [1, 2]] = np.nan
  726. self.frame.iloc[7, [0, 1]] = np.nan
  727. if axis == 0:
  728. frame = self.frame
  729. else:
  730. frame = self.frame.T
  731. grouped = frame.groupby(level=level, axis=axis, sort=sort)
  732. pieces = []
  733. def aggf(x):
  734. pieces.append(x)
  735. return getattr(x, op)(skipna=skipna, axis=axis)
  736. leftside = grouped.agg(aggf)
  737. rightside = getattr(frame, op)(level=level, axis=axis,
  738. skipna=skipna)
  739. if sort:
  740. rightside = rightside.sort_index(level=level, axis=axis)
  741. frame = frame.sort_index(level=level, axis=axis)
  742. # for good measure, groupby detail
  743. level_index = frame._get_axis(axis).levels[level]
  744. tm.assert_index_equal(leftside._get_axis(axis), level_index)
  745. tm.assert_index_equal(rightside._get_axis(axis), level_index)
  746. tm.assert_frame_equal(leftside, rightside)
  747. def test_stat_op_corner(self):
  748. obj = Series([10.0], index=MultiIndex.from_tuples([(2, 3)]))
  749. result = obj.sum(level=0)
  750. expected = Series([10.0], index=[2])
  751. tm.assert_series_equal(result, expected)
  752. def test_frame_any_all_group(self):
  753. df = DataFrame(
  754. {'data': [False, False, True, False, True, False, True]},
  755. index=[
  756. ['one', 'one', 'two', 'one', 'two', 'two', 'two'],
  757. [0, 1, 0, 2, 1, 2, 3]])
  758. result = df.any(level=0)
  759. ex = DataFrame({'data': [False, True]}, index=['one', 'two'])
  760. tm.assert_frame_equal(result, ex)
  761. result = df.all(level=0)
  762. ex = DataFrame({'data': [False, False]}, index=['one', 'two'])
  763. tm.assert_frame_equal(result, ex)
  764. def test_std_var_pass_ddof(self):
  765. index = MultiIndex.from_arrays([np.arange(5).repeat(10), np.tile(
  766. np.arange(10), 5)])
  767. df = DataFrame(np.random.randn(len(index), 5), index=index)
  768. for meth in ['var', 'std']:
  769. ddof = 4
  770. alt = lambda x: getattr(x, meth)(ddof=ddof)
  771. result = getattr(df[0], meth)(level=0, ddof=ddof)
  772. expected = df[0].groupby(level=0).agg(alt)
  773. tm.assert_series_equal(result, expected)
  774. result = getattr(df, meth)(level=0, ddof=ddof)
  775. expected = df.groupby(level=0).agg(alt)
  776. tm.assert_frame_equal(result, expected)
  777. def test_frame_series_agg_multiple_levels(self):
  778. result = self.ymd.sum(level=['year', 'month'])
  779. expected = self.ymd.groupby(level=['year', 'month']).sum()
  780. tm.assert_frame_equal(result, expected)
  781. result = self.ymd['A'].sum(level=['year', 'month'])
  782. expected = self.ymd['A'].groupby(level=['year', 'month']).sum()
  783. tm.assert_series_equal(result, expected)
  784. def test_groupby_multilevel(self):
  785. result = self.ymd.groupby(level=[0, 1]).mean()
  786. k1 = self.ymd.index.get_level_values(0)
  787. k2 = self.ymd.index.get_level_values(1)
  788. expected = self.ymd.groupby([k1, k2]).mean()
  789. # TODO groupby with level_values drops names
  790. tm.assert_frame_equal(result, expected, check_names=False)
  791. assert result.index.names == self.ymd.index.names[:2]
  792. result2 = self.ymd.groupby(level=self.ymd.index.names[:2]).mean()
  793. tm.assert_frame_equal(result, result2)
  794. def test_groupby_multilevel_with_transform(self):
  795. pass
  796. def test_multilevel_consolidate(self):
  797. index = MultiIndex.from_tuples([('foo', 'one'), ('foo', 'two'), (
  798. 'bar', 'one'), ('bar', 'two')])
  799. df = DataFrame(np.random.randn(4, 4), index=index, columns=index)
  800. df['Totals', ''] = df.sum(1)
  801. df = df._consolidate()
  802. def test_ix_preserve_names(self):
  803. result = self.ymd.loc[2000]
  804. result2 = self.ymd['A'].loc[2000]
  805. assert result.index.names == self.ymd.index.names[1:]
  806. assert result2.index.names == self.ymd.index.names[1:]
  807. result = self.ymd.loc[2000, 2]
  808. result2 = self.ymd['A'].loc[2000, 2]
  809. assert result.index.name == self.ymd.index.names[2]
  810. assert result2.index.name == self.ymd.index.names[2]
  811. def test_unstack_preserve_types(self):
  812. # GH #403
  813. self.ymd['E'] = 'foo'
  814. self.ymd['F'] = 2
  815. unstacked = self.ymd.unstack('month')
  816. assert unstacked['A', 1].dtype == np.float64
  817. assert unstacked['E', 1].dtype == np.object_
  818. assert unstacked['F', 1].dtype == np.float64
  819. def test_unstack_group_index_overflow(self):
  820. codes = np.tile(np.arange(500), 2)
  821. level = np.arange(500)
  822. index = MultiIndex(levels=[level] * 8 + [[0, 1]],
  823. codes=[codes] * 8 + [np.arange(2).repeat(500)])
  824. s = Series(np.arange(1000), index=index)
  825. result = s.unstack()
  826. assert result.shape == (500, 2)
  827. # test roundtrip
  828. stacked = result.stack()
  829. tm.assert_series_equal(s, stacked.reindex(s.index))
  830. # put it at beginning
  831. index = MultiIndex(levels=[[0, 1]] + [level] * 8,
  832. codes=[np.arange(2).repeat(500)] + [codes] * 8)
  833. s = Series(np.arange(1000), index=index)
  834. result = s.unstack(0)
  835. assert result.shape == (500, 2)
  836. # put it in middle
  837. index = MultiIndex(levels=[level] * 4 + [[0, 1]] + [level] * 4,
  838. codes=([codes] * 4 + [np.arange(2).repeat(500)] +
  839. [codes] * 4))
  840. s = Series(np.arange(1000), index=index)
  841. result = s.unstack(4)
  842. assert result.shape == (500, 2)
  843. def test_pyint_engine(self):
  844. # GH 18519 : when combinations of codes cannot be represented in 64
  845. # bits, the index underlying the MultiIndex engine works with Python
  846. # integers, rather than uint64.
  847. N = 5
  848. keys = [tuple(l) for l in [[0] * 10 * N,
  849. [1] * 10 * N,
  850. [2] * 10 * N,
  851. [np.nan] * N + [2] * 9 * N,
  852. [0] * N + [2] * 9 * N,
  853. [np.nan] * N + [2] * 8 * N + [0] * N]]
  854. # Each level contains 4 elements (including NaN), so it is represented
  855. # in 2 bits, for a total of 2*N*10 = 100 > 64 bits. If we were using a
  856. # 64 bit engine and truncating the first levels, the fourth and fifth
  857. # keys would collide; if truncating the last levels, the fifth and
  858. # sixth; if rotating bits rather than shifting, the third and fifth.
  859. for idx in range(len(keys)):
  860. index = MultiIndex.from_tuples(keys)
  861. assert index.get_loc(keys[idx]) == idx
  862. expected = np.arange(idx + 1, dtype=np.intp)
  863. result = index.get_indexer([keys[i] for i in expected])
  864. tm.assert_numpy_array_equal(result, expected)
  865. # With missing key:
  866. idces = range(len(keys))
  867. expected = np.array([-1] + list(idces), dtype=np.intp)
  868. missing = tuple([0, 1] * 5 * N)
  869. result = index.get_indexer([missing] + [keys[i] for i in idces])
  870. tm.assert_numpy_array_equal(result, expected)
  871. def test_to_html(self):
  872. self.ymd.columns.name = 'foo'
  873. self.ymd.to_html()
  874. self.ymd.T.to_html()
  875. def test_level_with_tuples(self):
  876. index = MultiIndex(levels=[[('foo', 'bar', 0), ('foo', 'baz', 0), (
  877. 'foo', 'qux', 0)], [0, 1]],
  878. codes=[[0, 0, 1, 1, 2, 2], [0, 1, 0, 1, 0, 1]])
  879. series = Series(np.random.randn(6), index=index)
  880. frame = DataFrame(np.random.randn(6, 4), index=index)
  881. result = series[('foo', 'bar', 0)]
  882. result2 = series.loc[('foo', 'bar', 0)]
  883. expected = series[:2]
  884. expected.index = expected.index.droplevel(0)
  885. tm.assert_series_equal(result, expected)
  886. tm.assert_series_equal(result2, expected)
  887. with pytest.raises(KeyError, match=r"^\(\('foo', 'bar', 0\), 2\)$"):
  888. series[('foo', 'bar', 0), 2]
  889. result = frame.loc[('foo', 'bar', 0)]
  890. result2 = frame.xs(('foo', 'bar', 0))
  891. expected = frame[:2]
  892. expected.index = expected.index.droplevel(0)
  893. tm.assert_frame_equal(result, expected)
  894. tm.assert_frame_equal(result2, expected)
  895. index = MultiIndex(levels=[[('foo', 'bar'), ('foo', 'baz'), (
  896. 'foo', 'qux')], [0, 1]],
  897. codes=[[0, 0, 1, 1, 2, 2], [0, 1, 0, 1, 0, 1]])
  898. series = Series(np.random.randn(6), index=index)
  899. frame = DataFrame(np.random.randn(6, 4), index=index)
  900. result = series[('foo', 'bar')]
  901. result2 = series.loc[('foo', 'bar')]
  902. expected = series[:2]
  903. expected.index = expected.index.droplevel(0)
  904. tm.assert_series_equal(result, expected)
  905. tm.assert_series_equal(result2, expected)
  906. result = frame.loc[('foo', 'bar')]
  907. result2 = frame.xs(('foo', 'bar'))
  908. expected = frame[:2]
  909. expected.index = expected.index.droplevel(0)
  910. tm.assert_frame_equal(result, expected)
  911. tm.assert_frame_equal(result2, expected)
  912. def test_mixed_depth_drop(self):
  913. arrays = [['a', 'top', 'top', 'routine1', 'routine1', 'routine2'],
  914. ['', 'OD', 'OD', 'result1', 'result2', 'result1'],
  915. ['', 'wx', 'wy', '', '', '']]
  916. tuples = sorted(zip(*arrays))
  917. index = MultiIndex.from_tuples(tuples)
  918. df = DataFrame(randn(4, 6), columns=index)
  919. result = df.drop('a', axis=1)
  920. expected = df.drop([('a', '', '')], axis=1)
  921. tm.assert_frame_equal(expected, result)
  922. result = df.drop(['top'], axis=1)
  923. expected = df.drop([('top', 'OD', 'wx')], axis=1)
  924. expected = expected.drop([('top', 'OD', 'wy')], axis=1)
  925. tm.assert_frame_equal(expected, result)
  926. result = df.drop(('top', 'OD', 'wx'), axis=1)
  927. expected = df.drop([('top', 'OD', 'wx')], axis=1)
  928. tm.assert_frame_equal(expected, result)
  929. expected = df.drop([('top', 'OD', 'wy')], axis=1)
  930. expected = df.drop('top', axis=1)
  931. result = df.drop('result1', level=1, axis=1)
  932. expected = df.drop([('routine1', 'result1', ''),
  933. ('routine2', 'result1', '')], axis=1)
  934. tm.assert_frame_equal(expected, result)
  935. def test_drop_nonunique(self):
  936. df = DataFrame([["x-a", "x", "a", 1.5], ["x-a", "x", "a", 1.2],
  937. ["z-c", "z", "c", 3.1], ["x-a", "x", "a", 4.1],
  938. ["x-b", "x", "b", 5.1], ["x-b", "x", "b", 4.1],
  939. ["x-b", "x", "b", 2.2],
  940. ["y-a", "y", "a", 1.2], ["z-b", "z", "b", 2.1]],
  941. columns=["var1", "var2", "var3", "var4"])
  942. grp_size = df.groupby("var1").size()
  943. drop_idx = grp_size.loc[grp_size == 1]
  944. idf = df.set_index(["var1", "var2", "var3"])
  945. # it works! #2101
  946. result = idf.drop(drop_idx.index, level=0).reset_index()
  947. expected = df[-df.var1.isin(drop_idx.index)]
  948. result.index = expected.index
  949. tm.assert_frame_equal(result, expected)
  950. def test_mixed_depth_pop(self):
  951. arrays = [['a', 'top', 'top', 'routine1', 'routine1', 'routine2'],
  952. ['', 'OD', 'OD', 'result1', 'result2', 'result1'],
  953. ['', 'wx', 'wy', '', '', '']]
  954. tuples = sorted(zip(*arrays))
  955. index = MultiIndex.from_tuples(tuples)
  956. df = DataFrame(randn(4, 6), columns=index)
  957. df1 = df.copy()
  958. df2 = df.copy()
  959. result = df1.pop('a')
  960. expected = df2.pop(('a', '', ''))
  961. tm.assert_series_equal(expected, result, check_names=False)
  962. tm.assert_frame_equal(df1, df2)
  963. assert result.name == 'a'
  964. expected = df1['top']
  965. df1 = df1.drop(['top'], axis=1)
  966. result = df2.pop('top')
  967. tm.assert_frame_equal(expected, result)
  968. tm.assert_frame_equal(df1, df2)
  969. def test_reindex_level_partial_selection(self):
  970. result = self.frame.reindex(['foo', 'qux'], level=0)
  971. expected = self.frame.iloc[[0, 1, 2, 7, 8, 9]]
  972. tm.assert_frame_equal(result, expected)
  973. result = self.frame.T.reindex(['foo', 'qux'], axis=1, level=0)
  974. tm.assert_frame_equal(result, expected.T)
  975. result = self.frame.loc[['foo', 'qux']]
  976. tm.assert_frame_equal(result, expected)
  977. result = self.frame['A'].loc[['foo', 'qux']]
  978. tm.assert_series_equal(result, expected['A'])
  979. result = self.frame.T.loc[:, ['foo', 'qux']]
  980. tm.assert_frame_equal(result, expected.T)
  981. def test_drop_level(self):
  982. result = self.frame.drop(['bar', 'qux'], level='first')
  983. expected = self.frame.iloc[[0, 1, 2, 5, 6]]
  984. tm.assert_frame_equal(result, expected)
  985. result = self.frame.drop(['two'], level='second')
  986. expected = self.frame.iloc[[0, 2, 3, 6, 7, 9]]
  987. tm.assert_frame_equal(result, expected)
  988. result = self.frame.T.drop(['bar', 'qux'], axis=1, level='first')
  989. expected = self.frame.iloc[[0, 1, 2, 5, 6]].T
  990. tm.assert_frame_equal(result, expected)
  991. result = self.frame.T.drop(['two'], axis=1, level='second')
  992. expected = self.frame.iloc[[0, 2, 3, 6, 7, 9]].T
  993. tm.assert_frame_equal(result, expected)
  994. def test_drop_level_nonunique_datetime(self):
  995. # GH 12701
  996. idx = Index([2, 3, 4, 4, 5], name='id')
  997. idxdt = pd.to_datetime(['201603231400',
  998. '201603231500',
  999. '201603231600',
  1000. '201603231600',
  1001. '201603231700'])
  1002. df = DataFrame(np.arange(10).reshape(5, 2),
  1003. columns=list('ab'), index=idx)
  1004. df['tstamp'] = idxdt
  1005. df = df.set_index('tstamp', append=True)
  1006. ts = Timestamp('201603231600')
  1007. assert df.index.is_unique is False
  1008. result = df.drop(ts, level='tstamp')
  1009. expected = df.loc[idx != 4]
  1010. tm.assert_frame_equal(result, expected)
  1011. @pytest.mark.parametrize('box', [Series, DataFrame])
  1012. def test_drop_tz_aware_timestamp_across_dst(self, box):
  1013. # GH 21761
  1014. start = Timestamp('2017-10-29', tz='Europe/Berlin')
  1015. end = Timestamp('2017-10-29 04:00:00', tz='Europe/Berlin')
  1016. index = pd.date_range(start, end, freq='15min')
  1017. data = box(data=[1] * len(index), index=index)
  1018. result = data.drop(start)
  1019. expected_start = Timestamp('2017-10-29 00:15:00', tz='Europe/Berlin')
  1020. expected_idx = pd.date_range(expected_start, end, freq='15min')
  1021. expected = box(dat

Large files files are truncated, but you can click here to view the full file