Return a new variable with given set of dimensions. This method might be used to attach new dimension(s) to variable. When possible, this operation does not copy this variable's data. Parameters ---------- dim : str or sequence of str or dict Dim
(self, dim, shape=None)
| 1445 | |
| 1446 | @deprecate_dims |
| 1447 | def set_dims(self, dim, shape=None): |
| 1448 | """Return a new variable with given set of dimensions. |
| 1449 | This method might be used to attach new dimension(s) to variable. |
| 1450 | |
| 1451 | When possible, this operation does not copy this variable's data. |
| 1452 | |
| 1453 | Parameters |
| 1454 | ---------- |
| 1455 | dim : str or sequence of str or dict |
| 1456 | Dimensions to include on the new variable. If a dict, values are |
| 1457 | used to provide the sizes of new dimensions; otherwise, new |
| 1458 | dimensions are inserted with length 1. |
| 1459 | |
| 1460 | Returns |
| 1461 | ------- |
| 1462 | Variable |
| 1463 | """ |
| 1464 | if isinstance(dim, str): |
| 1465 | dim = [dim] |
| 1466 | |
| 1467 | if shape is None and is_dict_like(dim): |
| 1468 | shape = tuple(dim.values()) |
| 1469 | |
| 1470 | missing_dims = set(self.dims) - set(dim) |
| 1471 | if missing_dims: |
| 1472 | raise ValueError( |
| 1473 | f"new dimensions {dim!r} must be a superset of " |
| 1474 | f"existing dimensions {self.dims!r}" |
| 1475 | ) |
| 1476 | |
| 1477 | self_dims = set(self.dims) |
| 1478 | expanded_dims = tuple(d for d in dim if d not in self_dims) + self.dims |
| 1479 | |
| 1480 | if self.dims == expanded_dims: |
| 1481 | # don't use broadcast_to unless necessary so the result remains |
| 1482 | # writeable if possible |
| 1483 | expanded_data = self._data |
| 1484 | elif shape is None or all( |
| 1485 | s == 1 for s, e in zip(shape, dim, strict=True) if e not in self_dims |
| 1486 | ): |
| 1487 | # "Trivial" broadcasting, i.e. simply inserting a new dimension |
| 1488 | # This is typically easier for duck arrays to implement |
| 1489 | # than the full "broadcast_to" semantics |
| 1490 | indexer = (None,) * (len(expanded_dims) - self.ndim) + (...,) |
| 1491 | # TODO: switch this to ._data once we teach ExplicitlyIndexed arrays to handle indexers with None. |
| 1492 | expanded_data = self.data[indexer] |
| 1493 | else: # elif shape is not None: |
| 1494 | dims_map = dict(zip(dim, shape, strict=True)) |
| 1495 | tmp_shape = tuple(dims_map[d] for d in expanded_dims) |
| 1496 | expanded_data = duck_array_ops.broadcast_to(self._data, tmp_shape) |
| 1497 | |
| 1498 | expanded_var = Variable( |
| 1499 | expanded_dims, expanded_data, self._attrs, self._encoding, fastpath=True |
| 1500 | ) |
| 1501 | return expanded_var.transpose(*dim) |
| 1502 | |
| 1503 | def _stack_once(self, dim: list[Hashable], new_dim: Hashable): |
| 1504 | if not set(dim) <= set(self.dims): |