(
self,
dims: Sequence[Hashable | EllipsisType],
new_dim: Hashable,
index_cls: type[Index],
create_index: bool | None = True,
)
| 5218 | return stack_index, stack_coords |
| 5219 | |
| 5220 | def _stack_once( |
| 5221 | self, |
| 5222 | dims: Sequence[Hashable | EllipsisType], |
| 5223 | new_dim: Hashable, |
| 5224 | index_cls: type[Index], |
| 5225 | create_index: bool | None = True, |
| 5226 | ) -> Self: |
| 5227 | if dims == ...: |
| 5228 | raise ValueError("Please use [...] for dims, rather than just ...") |
| 5229 | if ... in dims: |
| 5230 | dims = list(infix_dims(dims, self.dims)) |
| 5231 | |
| 5232 | new_variables: dict[Hashable, Variable] = {} |
| 5233 | stacked_var_names: list[Hashable] = [] |
| 5234 | drop_indexes: list[Hashable] = [] |
| 5235 | |
| 5236 | for name, var in self.variables.items(): |
| 5237 | if any(d in var.dims for d in dims): |
| 5238 | add_dims = [d for d in dims if d not in var.dims] |
| 5239 | vdims = list(var.dims) + add_dims |
| 5240 | shape = [self.sizes[d] for d in vdims] |
| 5241 | exp_var = var.set_dims(vdims, shape) |
| 5242 | stacked_var = exp_var.stack({new_dim: dims}) |
| 5243 | new_variables[name] = stacked_var |
| 5244 | stacked_var_names.append(name) |
| 5245 | else: |
| 5246 | new_variables[name] = var.copy(deep=False) |
| 5247 | |
| 5248 | # drop indexes of stacked coordinates (if any) |
| 5249 | for name in stacked_var_names: |
| 5250 | drop_indexes += list(self.xindexes.get_all_coords(name, errors="ignore")) |
| 5251 | |
| 5252 | new_indexes = {} |
| 5253 | new_coord_names = set(self._coord_names) |
| 5254 | if create_index or create_index is None: |
| 5255 | product_vars: dict[Any, Variable] = {} |
| 5256 | for dim in dims: |
| 5257 | idx, idx_vars = self._get_stack_index(dim, create_index=create_index) |
| 5258 | if idx is not None: |
| 5259 | product_vars.update(idx_vars) |
| 5260 | |
| 5261 | if len(product_vars) == len(dims): |
| 5262 | idx = index_cls.stack(product_vars, new_dim) |
| 5263 | new_indexes[new_dim] = idx |
| 5264 | new_indexes.update(dict.fromkeys(product_vars, idx)) |
| 5265 | idx_vars = idx.create_variables(product_vars) |
| 5266 | # keep consistent multi-index coordinate order |
| 5267 | for k in idx_vars: |
| 5268 | new_variables.pop(k, None) |
| 5269 | new_variables.update(idx_vars) |
| 5270 | new_coord_names.update(idx_vars) |
| 5271 | |
| 5272 | indexes = {k: v for k, v in self._indexes.items() if k not in drop_indexes} |
| 5273 | indexes.update(new_indexes) |
| 5274 | |
| 5275 | return self._replace_with_new_dims( |
| 5276 | new_variables, coord_names=new_coord_names, indexes=indexes |
| 5277 | ) |
no test coverage detected