Support flexible `weights` option in statistics.py module about esmlab

[andersy005]

Thinking about it a moment more, @matt-long, shouldn't we check for weights and x array shapes compatibility to make sure there isn't a mismatch before doing any computation too?

[matt-long]

@andersy005, do you have a proposed check? I think x must have all of the avg_over_dims_v (or equivalent) for the computation to be valid—or equivalently, x must have all of the weights.dims (but not the other way around).

[andersy005]

So far, I have this:

def _get_weights_and_dims(x, weights=None, dim=None):
    if dim and isinstance(dim, str):
        dims = [dim]

    elif isinstance(dim, list):
        dims = dim 

    
    else:
        dims = [k for k in x.dims]


    op_over_dims = [k for k in dims if k in x.dims]
    if not op_over_dims:
        raise ValueError("Unexpected dimensions for variable {0}".format(x.name))

    dims_shape = tuple(i for i in x[op_over_dims].dims.values())
    if weights is None:
        weights = xr.DataArray(np.ones(dims_shape), dims=op_over_dims)


    else:
        w = np.array(weights)
        assert w.shape == dims_shape
        weights = xr.DataArray(w, dims=op_over_dims)


    return weights, op_over_dims

def weighted_sum(x, weights=None, dim=None):

    if weights is None:
        warn("Computing sum with equal weights for all data points")

    
    weights, op_over_dims = _get_weights_and_dims(x, weights, dim)
    x_w_sum = (x * weights).sum(op_over_dims)

    original_attrs, original_encoding = get_original_attrs(x)
    return update_attrs(x_w_sum, original_attrs, original_encoding)

This line

dims_shape = tuple(i for i in x[op_over_dims].dims.values())

isn't working on dataarray. It works on datasets only.

[andersy005]

Let me know if you have a solution for me :)

[matt-long]

So I think you can use something like this.

dims_shape = tuple(l for i, l in enumerate(x.shape) if x.dims[i] in op_over_dims)

I would suggest changing this

...
else:
        w = np.array(weights)
        assert w.shape == dims_shape
        weights = xr.DataArray(w, dims=op_over_dims)

to

...
else:
     assert weights.shape == dims_shape

I don't see any reason to convert to numpy and then back.

We have to be aware of an issue with this, however. We have the following

def _apply_nan_mask(x, weights, avg_over_dims_v):
    weights = weights.where(x.notnull())
    ...

which is used in weighted_mean etc. This step effectively adds all the dimensions in x that are not already in weights to weights. I think this has desirable effects; for instance, in averaging data with non-uniform missing values in time, we are doing the right thing.

One possibility is to fold the _apply_nan_mask into your new _get_weights_and_dims function so we can ensure that the masking happens last.