Repository: dask/dask

Issue:
Weighted mean of masked dask array does not apply mask to weights
```python
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
https://github.com/dask/dask/blob/e823ecf51032ea7e4d57687a91d788cfb28e36ac/dask/array/routines.py#L1293

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.


Discussion:
@dkillick

@jcrist according to @bjlittle
I agree that @jcrist is likely the person who best knows the situation here.  However, if the fix is feasible and not too crazy then I suggest that @jcrist helps someone in the Met office implement the fix, rather than fixing it directly.  I think that it would be good to start transferring knowledge over to you all if you're up for some work.
So the above example is using `np.ma.average` (a masked implementation), but `da.average` (not a masked implementation). Looks like all that's needed is a masked port. I'd handle this by doing the following:

- Look at how `np.ma.average` differs from `np.average`
- Port over `np.ma.average` to use dask functions (for `da.average` it was mostly just replacing numpy functions with their dask equivalents)
- See what shared functionality there can be between `da.average` and `da.ma.average`. On first glance it looks to be as simple as masking the weights and then forwarding to `da.average`.

I agree with @mrocklin that this would be a good first issue if anyone wants to take a try at it.
@mrocklin I think a number of us here are keen to know dask as more than users 👍 
> See what shared functionality there can be between `da.average` and `da.ma.average`. On first glance it looks to be as simple as masking the weights and then forwarding to `da.average`.

After looking at all three functions I agree that there does not seems to happen any more magic in the masked version and passing the masked weights should be all that is required. Actually I am somewhat surprised that numpy does not do this internally as well...

Resolving patch:
diff --git a/dask/array/ma.py b/dask/array/ma.py
--- a/dask/array/ma.py
+++ b/dask/array/ma.py
@@ -8,6 +8,7 @@
 from ..base import normalize_token
 from .core import (concatenate_lookup, tensordot_lookup, map_blocks,
                    asanyarray, atop)
+from .routines import _average
 
 
 if LooseVersion(np.__version__) < '1.11.2':
@@ -249,3 +250,8 @@ def set_fill_value(a, fill_value):
     res = a.map_blocks(_set_fill_value, fill_value)
     a.dask = res.dask
     a.name = res.name
+
+
+@wraps(np.ma.average)
+def average(a, axis=None, weights=None, returned=False):
+    return _average(a, axis, weights, returned, is_masked=True)
diff --git a/dask/array/routines.py b/dask/array/routines.py
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -1292,11 +1292,11 @@ def insert(arr, obj, values, axis):
     return concatenate(interleaved, axis=axis)
 
 
-@wraps(np.average)
-def average(a, axis=None, weights=None, returned=False):
+def _average(a, axis=None, weights=None, returned=False, is_masked=False):
     # This was minimally modified from numpy.average
     # See numpy license at https://github.com/numpy/numpy/blob/master/LICENSE.txt
     # or NUMPY_LICENSE.txt within this directory
+    # Wrapper used by da.average or da.ma.average.
     a = asanyarray(a)
 
     if weights is None:
@@ -1326,7 +1326,9 @@ def average(a, axis=None, weights=None, returned=False):
             # setup wgt to broadcast along axis
             wgt = broadcast_to(wgt, (a.ndim - 1) * (1,) + wgt.shape)
             wgt = wgt.swapaxes(-1, axis)
-
+        if is_masked:
+            from .ma import getmaskarray
+            wgt = wgt * (~getmaskarray(a))
         scl = wgt.sum(axis=axis, dtype=result_dtype)
         avg = multiply(a, wgt, dtype=result_dtype).sum(axis) / scl
 
@@ -1336,3 +1338,8 @@ def average(a, axis=None, weights=None, returned=False):
         return avg, scl
     else:
         return avg
+
+
+@wraps(np.average)
+def average(a, axis=None, weights=None, returned=False):
+    return _average(a, axis, weights, returned, is_masked=False)


Test patch:
diff --git a/dask/array/tests/test_masked.py b/dask/array/tests/test_masked.py
--- a/dask/array/tests/test_masked.py
+++ b/dask/array/tests/test_masked.py
@@ -331,3 +331,20 @@ def test_set_fill_value():
 
     with pytest.raises(ValueError):
         da.ma.set_fill_value(dmx, dx)
+
+
+def test_average_weights_with_masked_array():
+    mask = np.array([[True, False],
+                     [True, True],
+                     [False, True]])
+    data = np.arange(6).reshape((3, 2))
+    a = np.ma.array(data, mask=mask)
+    d_a = da.ma.masked_array(data=data, mask=mask, chunks=2)
+
+    weights = np.array([0.25, 0.75])
+    d_weights = da.from_array(weights, chunks=2)
+
+    np_avg = np.ma.average(a, weights=weights, axis=1)
+    da_avg = da.ma.average(d_a, weights=d_weights, axis=1)
+
+    assert_eq(np_avg, da_avg)