Function apply_groupby_func

xarray/computation/apply_ufunc.py:574–618

Apply a dataset or datarray level function over GroupBy, Dataset, DataArray, Variable and/or ndarray objects.

(func, *args)

def apply_groupby_func(func, *args):
    """Apply a dataset or datarray level function over GroupBy, Dataset,
    DataArray, Variable and/or ndarray objects.
    """
    from xarray.core.groupby import GroupBy, peek_at

    groupbys = [arg for arg in args if isinstance(arg, GroupBy)]
    assert groupbys, "must have at least one groupby to iterate over"
    first_groupby = groupbys[0]
    (grouper,) = first_groupby.groupers
    if any(not grouper.group.equals(gb.groupers[0].group) for gb in groupbys[1:]):  # type: ignore[union-attr]
        raise ValueError(
            "apply_ufunc can only perform operations over "
            "multiple GroupBy objects at once if they are all "
            "grouped the same way"
        )

    grouped_dim = grouper.name
    unique_values = grouper.unique_coord.values

    iterators = []
    for arg in args:
        iterator: Iterator[Any]
        if isinstance(arg, GroupBy):
            iterator = (value for _, value in arg)
        elif hasattr(arg, "dims") and grouped_dim in arg.dims:
            if isinstance(arg, Variable):
                raise ValueError(
                    "groupby operations cannot be performed with "
                    "xarray.Variable objects that share a dimension with "
                    "the grouped dimension"
                )
            iterator = _iter_over_selections(arg, grouped_dim, unique_values)
        else:
            iterator = itertools.repeat(arg)
        iterators.append(iterator)

    applied: Iterator = itertools.starmap(func, zip(*iterators, strict=False))
    applied_example, applied = peek_at(applied)
    combine = first_groupby._combine  # type: ignore[attr-defined]
    if isinstance(applied_example, tuple):
        combined = tuple(combine(output) for output in zip(*applied, strict=True))
    else:
        combined = combine(applied)
    return combined


def unified_dim_sizes(