weighted operations: performance optimisations

[mathause]

There was a discussion on the performance of the weighted mean/ sum in terms of memory footprint but also speed, and there may indeed be some things that can be optimized. See the posts at the end of the PR. However, the optimal implementation will probably depend on the use case and some profiling will be required.

I'll just open an issue to keep track of this.
@seth-p

[mathause]

maybe relevant: #1995

[mathause]

%load_ext line_profiler


import numpy as np
import xarray as xr

from xarray.core.weighted import Weighted as w

shape_weights = (1000, 1000)
shape_data = (1000, 1000, 10)
add_nans = False

def lprun_weighted(shape_weights, shape_data, add_nans, skipna=None):

    weights = xr.DataArray(np.random.randn(*shape_weights))

    data = np.random.randn(*shape_data)

    # add approximately 25 % NaNs
    if add_nans:
        c = int(data.size * 0.25)
        data.ravel()[np.random.choice(data.size, c, replace=False)] = np.NaN

    data = xr.DataArray(data)


    return data.weighted(weights).mean(skipna=skipna)



%lprun -f w._reduce -f w._weighted_mean -f w._sum_of_weights -f w._weighted_sum -f w.__init__ -f lprun_weighted -u 1e-03 lprun_weighted(shape_weights, shape_data, add_nans, skipna=None)

[mathause]

weighted(weights).mean(skipna=True) calls

xarray/xarray/core/weighted.py

Line 143 in d1f7cb8

mask = da.notnull()

and

xarray/xarray/core/weighted.py

Line 130 in d1f7cb8

da = da.fillna(0.0)

da.fillna(0.0) in turn calls where(null(data), other). Thus null/ notnull is called twice. This could be optimized by doing self.null = obj.null() in __init__(...). This might also allow self.any_null = self.null.any() to skip NaN handling if there are none. However, this needs more thinking if a Dataset is passed. Probably overkill but leaving this here for reference.