Weighted mean of masked dask array does not apply mask to weights #3846
Description
import dask.array as da
import numpy as np
my_arr = np.array(range(10))
my_mask = [True, True, True, False, False, False, False, False, False, False]
my_weights = np.array([100, 100, 2, 2, 1, 1, 1, 1, 1, 1], dtype='float64')
my_masked_arr = np.ma.array(my_arr, mask=my_mask)
print("numpy weighted masked avg (unmasked weights): ", ma.average(my_masked_arr, weights=my_weights))
my_masked_dask_arr = da.ma.masked_array(data=my_arr, mask=my_mask)
print("dask weights masked avg (unmasked weights): ", da.average(my_masked_dask_arr, weights=my_weights).compute())
my_masked_weights = ma.masked_array(my_weights, my_mask)
print("dask weights masked avg (weights): ", da.average(my_masked_dask_arr, weights=my_masked_weights).compute())gives this output
numpy weighted masked avg (unmasked weights): 5.625
dask weights masked avg (unmasked weights): 0.214285714286
dask weights masked avg (weights): 5.625
Upon further investigation it seems that the weights are normalised by dask, but that normalisation doesn't take into account the array mask.
i.e. this line
Line 1293 in e823ecf
I'm not sure how to fix this - I don't think one has access to the mask without realising all the data, so the weight normalisation can't happen lazily(?) Maybe the way to solve this is to specify that the weights array must be masked by the user before being passed - this would make the API different to numpy.