Class UniqueGrouper

xarray/groupers.py:210–340

Grouper object for grouping by a categorical variable. Parameters ---------- labels: array-like, optional Group labels to aggregate on. This is required when grouping by a chunked array type (e.g. dask or cubed) since it is used to construct the coordinate on the ou

@dataclass
class UniqueGrouper(Grouper):
    """
    Grouper object for grouping by a categorical variable.

    Parameters
    ----------
    labels: array-like, optional
        Group labels to aggregate on. This is required when grouping by a chunked array type
        (e.g. dask or cubed) since it is used to construct the coordinate on the output.
        Grouped operations will only be run on the specified group labels. Any group that is not
        present in ``labels`` will be ignored.
    """

    _group_as_index: pd.Index | None = field(default=None, repr=False, init=False)
    labels: ArrayLike | None = field(default=None)

    @property
    def group_as_index(self) -> pd.Index:
        """Caches the group DataArray as a pandas Index."""
        if self._group_as_index is None:
            if self.group.ndim == 1:
                self._group_as_index = self.group.to_index()
            else:
                self._group_as_index = pd.Index(np.array(self.group).ravel())
        return self._group_as_index

    def reset(self) -> Self:
        return type(self)()

    def factorize(self, group: T_Group) -> EncodedGroups:
        self.group = group

        if is_chunked_array(group.data) and self.labels is None:
            raise ValueError(
                "When grouping by a dask array, `labels` must be passed using "
                "a UniqueGrouper object."
            )
        if self.labels is not None:
            return self._factorize_given_labels(group)

        index = self.group_as_index
        is_unique_and_monotonic = isinstance(self.group, _DummyGroup) or (
            index.is_unique
            and (index.is_monotonic_increasing or index.is_monotonic_decreasing)
        )
        is_dimension = self.group.dims == (self.group.name,)
        can_squeeze = is_dimension and is_unique_and_monotonic

        if can_squeeze:
            return self._factorize_dummy()
        else:
            return self._factorize_unique()

    def _factorize_given_labels(self, group: T_Group) -> EncodedGroups:
        codes = apply_ufunc(
            _factorize_given_labels,
            group,
            kwargs={"labels": self.labels},