|
22 | 22 | import pandas.core.common as common
|
23 | 23 | import pandas._tseries as _tseries
|
24 | 24 |
|
| 25 | + |
| 26 | +def _ensure_like_indices(time, panels): |
| 27 | + """ |
| 28 | + Makes sure that time and panels are conformable |
| 29 | + """ |
| 30 | + n_time = len(time) |
| 31 | + n_panel = len(panels) |
| 32 | + u_panels = np.unique(panels) # this sorts! |
| 33 | + u_time = np.unique(time) |
| 34 | + if len(u_time) == n_time: |
| 35 | + time = np.tile(u_time, len(u_panels)) |
| 36 | + if len(u_panels) == n_panel: |
| 37 | + panels = np.repeat(u_panels, len(u_time)) |
| 38 | + return time, panels |
| 39 | + |
| 40 | +def panel_index(time, panels, names=['time', 'panel']): |
| 41 | + """ |
| 42 | + Returns a multi-index suitable for a panel-like DataFrame |
| 43 | +
|
| 44 | + Parameters |
| 45 | + ---------- |
| 46 | + time : array-like |
| 47 | + Time index, does not have to repeat |
| 48 | + panels : array-like |
| 49 | + Panel index, does not have to repeat |
| 50 | + names : list, optional |
| 51 | + List containing the names of the indices |
| 52 | +
|
| 53 | + Returns |
| 54 | + ------- |
| 55 | + multi_index : MultiIndex |
| 56 | + Time index is the first level, the panels are the second level. |
| 57 | +
|
| 58 | + Examples |
| 59 | + -------- |
| 60 | + >>> years = range(1960,1963) |
| 61 | + >>> panels = ['A', 'B', 'C'] |
| 62 | + >>> panel_idx = panel_index(years, panels) |
| 63 | + >>> panel_idx |
| 64 | + MultiIndex([(1960, 'A'), (1961, 'A'), (1962, 'A'), (1960, 'B'), (1961, 'B'), |
| 65 | + (1962, 'B'), (1960, 'C'), (1961, 'C'), (1962, 'C')], dtype=object) |
| 66 | + |
| 67 | + or |
| 68 | +
|
| 69 | + >>> import numpy as np |
| 70 | + >>> years = np.repeat(range(1960,1963), 3) |
| 71 | + >>> panels = np.tile(['A', 'B', 'C'], 3) |
| 72 | + >>> panel_idx = panel_index(years, panels) |
| 73 | + >>> panel_idx |
| 74 | + MultiIndex([(1960, 'A'), (1960, 'B'), (1960, 'C'), (1961, 'A'), (1961, 'B'), |
| 75 | + (1961, 'C'), (1962, 'A'), (1962, 'B'), (1962, 'C')], dtype=object) |
| 76 | + """ |
| 77 | + time, panels = _ensure_like_indices(time, panels) |
| 78 | + time_factor = Factor(time) |
| 79 | + panel_factor = Factor(panels) |
| 80 | + |
| 81 | + labels = [time_factor.labels, panel_factor.labels] |
| 82 | + levels = [time_factor.levels, panel_factor.levels] |
| 83 | + return MultiIndex(levels, labels, sortorder=None, names=names) |
| 84 | + |
25 | 85 | class PanelError(Exception):
|
26 | 86 | pass
|
27 | 87 |
|
|
0 commit comments