Create a new multi-index maybe by expanding an existing one with new variables as index levels. The index and its corresponding coordinates may be created along a new dimension.
(
cls,
dim: Hashable,
current_variables: Mapping[Any, Variable],
variables: Mapping[Any, Variable],
)
| 1164 | |
| 1165 | @classmethod |
| 1166 | def from_variables_maybe_expand( |
| 1167 | cls, |
| 1168 | dim: Hashable, |
| 1169 | current_variables: Mapping[Any, Variable], |
| 1170 | variables: Mapping[Any, Variable], |
| 1171 | ) -> tuple[PandasMultiIndex, IndexVars]: |
| 1172 | """Create a new multi-index maybe by expanding an existing one with |
| 1173 | new variables as index levels. |
| 1174 | |
| 1175 | The index and its corresponding coordinates may be created along a new dimension. |
| 1176 | """ |
| 1177 | names: list[Hashable] = [] |
| 1178 | codes: list[Iterable[int]] = [] |
| 1179 | levels: list[Iterable[Any]] = [] |
| 1180 | level_variables: dict[Any, Variable] = {} |
| 1181 | |
| 1182 | _check_dim_compat({**current_variables, **variables}) |
| 1183 | |
| 1184 | if len(current_variables) > 1: |
| 1185 | # expand from an existing multi-index |
| 1186 | data = cast( |
| 1187 | PandasMultiIndexingAdapter, next(iter(current_variables.values()))._data |
| 1188 | ) |
| 1189 | current_index = data.array |
| 1190 | names.extend(current_index.names) |
| 1191 | codes.extend(current_index.codes) |
| 1192 | levels.extend(current_index.levels) |
| 1193 | for name in current_index.names: |
| 1194 | level_variables[name] = current_variables[name] |
| 1195 | |
| 1196 | elif len(current_variables) == 1: |
| 1197 | # expand from one 1D variable (no multi-index): convert it to an index level |
| 1198 | var = next(iter(current_variables.values())) |
| 1199 | new_var_name = f"{dim}_level_0" |
| 1200 | names.append(new_var_name) |
| 1201 | cat = pd.Categorical(var.values, ordered=True) |
| 1202 | codes.append(cat.codes) |
| 1203 | levels.append(cat.categories) |
| 1204 | level_variables[new_var_name] = var |
| 1205 | |
| 1206 | for name, var in variables.items(): |
| 1207 | names.append(name) |
| 1208 | cat = pd.Categorical(var.values, ordered=True) |
| 1209 | codes.append(cat.codes) |
| 1210 | levels.append(cat.categories) |
| 1211 | level_variables[name] = var |
| 1212 | |
| 1213 | codes_as_lists = [list(x) for x in codes] |
| 1214 | levels_as_lists = [list(level) for level in levels] |
| 1215 | index = pd.MultiIndex(levels=levels_as_lists, codes=codes_as_lists, names=names) |
| 1216 | level_coords_dtype = {k: var.dtype for k, var in level_variables.items()} |
| 1217 | obj = cls(index, dim, level_coords_dtype=level_coords_dtype) |
| 1218 | index_vars = obj.create_variables(level_variables) |
| 1219 | |
| 1220 | return obj, index_vars |
| 1221 | |
| 1222 | def keep_levels( |
| 1223 | self, level_variables: Mapping[Any, Variable] |
no test coverage detected