Class ComposedGrouper

xarray/core/groupby.py:513–577

Helper class for multi-variable GroupBy. This satisfies the Grouper interface, but is awkward to wrap in ResolvedGrouper. For one, it simply re-infers a new EncodedGroups using known information in existing ResolvedGroupers. So passing in a `group` (hard to define), and `obj` (p

@dataclass
class ComposedGrouper:
    """
    Helper class for multi-variable GroupBy.
    This satisfies the Grouper interface, but is awkward to wrap in ResolvedGrouper.
    For one, it simply re-infers a new EncodedGroups using known information
    in existing ResolvedGroupers. So passing in a `group` (hard to define),
    and `obj` (pointless) is not useful.
    """

    groupers: tuple[ResolvedGrouper, ...]

    def factorize(self) -> EncodedGroups:
        from xarray.groupers import EncodedGroups

        groupers = self.groupers

        # At this point all arrays have been factorized.
        codes = tuple(grouper.codes for grouper in groupers)
        shape = tuple(grouper.size for grouper in groupers)
        masks = tuple((code == -1) for code in codes)
        # We broadcast the codes against each other
        broadcasted_codes = broadcast(*codes)
        # This fully broadcasted DataArray is used as a template later
        first_codes = broadcasted_codes[0]
        # Now we convert to a single variable GroupBy problem
        _flatcodes = np.ravel_multi_index(
            tuple(codes.data for codes in broadcasted_codes), shape, mode="wrap"
        )
        # NaNs; as well as values outside the bins are coded by -1
        # Restore these after the raveling
        broadcasted_masks = broadcast(*masks)
        mask = functools.reduce(np.logical_or, broadcasted_masks)  # type: ignore[arg-type]
        _flatcodes = where(mask.data, -1, _flatcodes)

        full_index = pd.MultiIndex.from_product(
            [list(grouper.full_index.values) for grouper in groupers],
            names=tuple(grouper.name for grouper in groupers),
        )
        if not full_index.is_unique:
            raise ValueError(
                "The output index for the GroupBy is non-unique. "
                "This is a bug in the Grouper provided."
            )
        # This will be unused when grouping by dask arrays, so skip..
        if not is_chunked_array(_flatcodes):
            # Constructing an index from the product is wrong when there are missing groups
            # (e.g. binning, resampling). Account for that now.
            midx = full_index[np.sort(pd.unique(_flatcodes[~mask]))]
            group_indices = _codes_to_group_indices(_flatcodes.ravel(), len(full_index))
        else:
            midx = full_index
            group_indices = None

        dim_name = "stacked_" + "_".join(str(grouper.name) for grouper in groupers)

        coords = Coordinates.from_pandas_multiindex(midx, dim=dim_name)
        for grouper in groupers:
            coords.variables[grouper.name].attrs = grouper.group.attrs