Ejemplo n.º 1
0
    def group(self, da: xr.DataArray = None, **das: xr.DataArray):
        """Return a xr.core.groupby.GroupBy object.

        More than one array can be combined to a dataset before grouping using the `das`  kwargs.
        A new `window` dimension is added if `self.window` is larger than 1.
        If `Grouper.dim` is 'time', but 'prop' is None, the whole array is grouped together.

        When multiple arrays are passed, some of them can be grouped along the same group as self.
        They are boadcasted, merged to the grouping dataset and regrouped in the output.
        """
        if das:
            from .utils import broadcast  # pylint: disable=cyclic-import

            if da is not None:
                das[da.name] = da

            da = xr.Dataset(
                data_vars={
                    name: das.pop(name)
                    for name in list(das.keys()) if self.dim in das[name].dims
                })

            # "Ungroup" the grouped arrays
            da = da.assign({
                name: broadcast(var,
                                da[self.dim],
                                group=self,
                                interp="nearest")
                for name, var in das.items()
            })

        if self.window > 1:
            da = da.rolling(center=True, **{
                self.dim: self.window
            }).construct(window_dim="window")
            if da.chunks is not None:
                # Rechunk. There might be padding chunks.
                da = da.chunk({self.dim: -1})

        if self.prop is None and self.dim == "time":
            group = self.get_index(da)
            group.name = self.dim
        else:
            group = self.name

        return da.groupby(group)