add average function

[mathause]

It would be nice to be able to do ds.average() to compute weighted averages (e.g. for geo data). Of course this would require the axes to be in a predictable order. Or to give a weight per dimension...

[shoyer]

Module error checking, etc., this would look something like:

def average(self, dim=None, weights=None):
    if weights is None:
        return self.mean(dim)
    else:
        return (self * weights).sum(dim) / weights.sum(dim)

This is pretty easy to do manually, but I can see the value in having the standard method around, so I'm definitely open to PRs to add this functionality.

[mathause]

This is has to be adjusted if there are NaN in the array. weights.sum(dim) needs to be corrected not to count weights on indices where there is a NaN in self.

Is there a better way to get the correct weights than:

total_weights = weights.sum(dim) * self / self

It should probably not be used on a Dataset as every DataArray may have its own NaN structure. Or the equivalent Dataset method should loop through the DataArrays.

[shoyer]

Possibly using where, e.g., weights.where(self.notnull()).sum(dim).

[mathause]

Thanks - that seems to be the fastest possibility. I wrote the functions for Dataset and DataArray

def average_da(self, dim=None, weights=None):
    """
    weighted average for DataArrays

    Parameters
    ----------
    dim : str or sequence of str, optional
        Dimension(s) over which to apply average.
    weights : DataArray
        weights to apply. Shape must be broadcastable to shape of self.

    Returns
    -------
    reduced : DataArray
        New DataArray with average applied to its data and the indicated
        dimension(s) removed.

    """

    if weights is None:
        return self.mean(dim)
    else:
        if not isinstance(weights, xray.DataArray):
            raise ValueError("weights must be a DataArray")

        # if NaNs are present, we need individual weights
        if self.notnull().any():
            total_weights = weights.where(self.notnull()).sum(dim=dim)
        else:
            total_weights = weights.sum(dim)

        return (self * weights).sum(dim) / total_weights

# -----------------------------------------------------------------------------

def average_ds(self, dim=None, weights=None):
    """
    weighted average for Datasets

    Parameters
    ----------
    dim : str or sequence of str, optional
        Dimension(s) over which to apply average.
    weights : DataArray
        weights to apply. Shape must be broadcastable to shape of data.

    Returns
    -------
    reduced : Dataset
        New Dataset with average applied to its data and the indicated
        dimension(s) removed.

    """

    if weights is None:
        return self.mean(dim)
    else:
        return self.apply(average_da, dim=dim, weights=weights)

They can be combined to one function:

def average(data, dim=None, weights=None):
    """
    weighted average for xray objects

    Parameters
    ----------
    data : Dataset or DataArray
        the xray object to average over
    dim : str or sequence of str, optional
        Dimension(s) over which to apply average.
    weights : DataArray
        weights to apply. Shape must be broadcastable to shape of data.

    Returns
    -------
    reduced : Dataset or DataArray
        New xray object with average applied to its data and the indicated
        dimension(s) removed.

    """

    if isinstance(data, xray.Dataset):
        return average_ds(data, dim, weights)
    elif isinstance(data, xray.DataArray):
        return average_da(data, dim, weights)
    else:
        raise ValueError("date must be an xray Dataset or DataArray")

Or a monkey patch:

xray.DataArray.average = average_da
xray.Dataset.average = average_ds

[jhamman]

@MaximilianR has suggested a groupby/rolling-like interface to weighted reductions.

da.weighted(weights=ds.dim).mean()
# or maybe
da.weighted(time=days_per_month(da.time)).mean()

I really like this idea, as does @shoyer. I'm going to close my PR in hopes of this becoming reality.

[shoyer]

I would suggest not using keyword arguments for weighted. Instead, just align based on the labels of the argument like regular xarray operations. So we'd write da.weighted(days_per_month(da.time)).mean()

[mathause]

Sounds like a clean solution. Then we can defer handling of NaN in the weights to weighted (e.g. by a skipna_weights argument in weighted). Also returning sum_of_weights can be a method of the class.

We may still end up implementing all required methods separately in weighted. For mean we do:

(data * weights / sum_of_weights).sum(dim=dim)

i.e. we use sum and not mean. We could rewrite this to:

(data * weights / sum_of_weights).mean(dim=dim) * weights.count(dim=dim)

However, I think this can not be generalized to a reduce function. See e.g. for std http://stackoverflow.com/questions/30383270/how-do-i-calculate-the-standard-deviation-between-weighted-measurements

Additionally, weighted does not make sense for many operations (I would say) e.g.: min, max, count, ...

[mathause]

Do we want

da.weighted(weight, dim='time').mean()

or

da.weighted(weight).mean(dim='time')

[jhamman]

@mathause -

I would think you want the latter (da.weighted(weight).mean(dim='time')). weighted should handle the brodcasting of weight such that you could do this:

>>> da.shape
(72, 10, 15)
>>> da.dims
('time', 'x', 'y')
>>> weights = some_func_of_time(time)
>>> da.weighted(weights).mean(dim=('time', 'x'))
...

[shoyer]

Yes, +1 for da.weighted(weight).mean(dim='time'). The mean method on weighted should have the same arguments as the mean method on DataArray -- it's just changed due to the context.

We may still end up implementing all required methods separately in weighted.

This is a fair point, I haven't looked in to the details of these implementations yet. But I expect there are still at least a few picks of logic that we will be able to share.

[pwolfram]

@mathause can you please comment on the status of this issue? Is there an associated PR somewhere? Thanks!