Repository: blaze/dask (python)

Issue:
dask.array.nanmax does not properly skip missing values
<img width="809" alt="screen shot 2015-10-06 at 9 44 14 pm" src="https://cloud.githubusercontent.com/assets/1217238/10329252/75538ca2-6c73-11e5-9afd-8474df535c31.png">

Still working on a minimal example with data I can share.

Pull request description:
BUG: fixes for nan-skipping aggregation functions
Fixes #773

This includes tests but no fix yet for `nanargmax` when it encounters an all-NaN slice in a chunk. I'll leave that to a follow up.

Patch:
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index ebce7c164..d75a2c541 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -149,14 +149,14 @@ with ignoring(AttributeError):
 
 @wraps(chunk.nanmin)
 def nanmin(a, axis=None, keepdims=False):
-    return reduction(a, chunk.nanmin, chunk.min, axis=axis, keepdims=keepdims,
+    return reduction(a, chunk.nanmin, chunk.nanmin, axis=axis, keepdims=keepdims,
                      dtype=a._dtype)
 
 
 @wraps(chunk.nanmax)
 def nanmax(a, axis=None, keepdims=False):
-    return reduction(a, chunk.nanmax, chunk.max, axis=axis, keepdims=keepdims,
-                     dtype=a._dtype)
+    return reduction(a, chunk.nanmax, chunk.nanmax, axis=axis,
+                     keepdims=keepdims, dtype=a._dtype)
 
 
 def numel(x, **kwargs):
@@ -226,7 +226,7 @@ def moment_chunk(A, order=2, sum=chunk.sum, numel=numel, dtype='f8', **kwargs):
     return result
 
 
-def moment_agg(data, order=2, ddof=0, dtype='f8', **kwargs):
+def moment_agg(data, order=2, ddof=0, dtype='f8', sum=np.sum, **kwargs):
     totals = data['total']
     ns = data['n']
     Ms = data['M']
@@ -237,7 +237,7 @@ def moment_agg(data, order=2, ddof=0, dtype='f8', **kwargs):
     keepdim_kw = kwargs.copy()
     keepdim_kw['keepdims'] = True
 
-    n = ns.sum(**keepdim_kw)
+    n = sum(ns, **keepdim_kw)
     mu = divide(totals.sum(**keepdim_kw), n, dtype=dtype)
     inner_term = divide(totals, ns, dtype=dtype) - mu
 
@@ -245,10 +245,10 @@ def moment_agg(data, order=2, ddof=0, dtype='f8', **kwargs):
 
     for k in range(1, order - 1):
         coeff = factorial(order)/(factorial(k)*factorial(order - k))
-        result += coeff * (Ms[..., order - k - 2] * inner_term**k).sum(**kwargs)
+        result += coeff * sum(Ms[..., order - k - 2] * inner_term**k, **kwargs)
 
-    result += (ns * inner_term**order).sum(**kwargs)
-    result = divide(result, (n.sum(**kwargs) - ddof), dtype=dtype)
+    result += sum(ns * inner_term**order, **kwargs)
+    result = divide(result, sum(n, **kwargs) - ddof, dtype=dtype)
     return result
 
 
@@ -286,7 +286,7 @@ def nanvar(a, axis=None, dtype=None, keepdims=False, ddof=0):
     else:
         dt = None
     return reduction(a, partial(moment_chunk, sum=chunk.nansum, numel=nannumel),
-                     partial(moment_agg, ddof=ddof), axis=axis,
+                     partial(moment_agg, sum=np.nansum, ddof=ddof), axis=axis,
                      keepdims=keepdims, dtype=dt)
 
 with ignoring(AttributeError):


Test patch:
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index ebb8ddb8f..018fb19f8 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -15,7 +15,7 @@ def eq(a, b):
     if isinstance(b, da.Array):
         b = b.compute()
     if isinstance(a, (np.generic, np.ndarray)):
-        return np.allclose(a, b)
+        return np.all(np.isclose(a, b, equal_nan=True))
     else:
         return a == b
 
@@ -187,6 +187,46 @@ def test_reductions_2D_int():
     assert eq(da.nanargmin(a, axis=1), np.nanargmin(x, axis=1))
 
 
+def test_reductions_2D_nans():
+    # chunks are a mix of some/all/no NaNs
+    x = np.full((4, 4), np.nan)
+    x[:2, :2] = np.array([[1, 2], [3, 4]])
+    x[2, 2] = 5
+    x[3, 3] = 6
+    a = da.from_array(x, chunks=(2, 2))
+
+    reduction_2d_test(da.sum, a, np.sum, x, False)
+    reduction_2d_test(da.prod, a, np.prod, x, False)
+    reduction_2d_test(da.mean, a, np.mean, x, False)
+    reduction_2d_test(da.var, a, np.var, x, False)
+    reduction_2d_test(da.std, a, np.std, x, False)
+    reduction_2d_test(da.min, a, np.min, x, False)
+    reduction_2d_test(da.max, a, np.max, x, False)
+    reduction_2d_test(da.any, a, np.any, x, False)
+    reduction_2d_test(da.all, a, np.all, x, False)
+
+    reduction_2d_test(da.nansum, a, np.nansum, x, False)
+    with ignoring(AttributeError):
+        reduction_2d_test(da.nanprod, a, np.nanprod, x, False)
+    reduction_2d_test(da.nanmean, a, np.nanmean, x, False)
+    reduction_2d_test(da.nanvar, a, np.nanvar, x, False)
+    reduction_2d_test(da.nanstd, a, np.nanstd, x, False)
+    reduction_2d_test(da.nanmin, a, np.nanmin, x, False)
+    reduction_2d_test(da.nanmax, a, np.nanmax, x, False)
+
+    # TODO: fix these tests, which fail with this error from NumPy:
+    # ValueError("All-NaN slice encountered"), because some of the chunks
+    # (not all) have all NaN values.
+    # assert eq(da.argmax(a, axis=0), np.argmax(x, axis=0))
+    # assert eq(da.argmin(a, axis=0), np.argmin(x, axis=0))
+    # assert eq(da.nanargmax(a, axis=0), np.nanargmax(x, axis=0))
+    # assert eq(da.nanargmin(a, axis=0), np.nanargmin(x, axis=0))
+    # assert eq(da.argmax(a, axis=1), np.argmax(x, axis=1))
+    # assert eq(da.argmin(a, axis=1), np.argmin(x, axis=1))
+    # assert eq(da.nanargmax(a, axis=1), np.nanargmax(x, axis=1))
+    # assert eq(da.nanargmin(a, axis=1), np.nanargmin(x, axis=1))
+
+
 def test_moment():
     def moment(x, n, axis=None):
         return ((x - x.mean(axis=axis, keepdims=True))**n).sum(