Repository: dask/dask

Issue:
da.nanargmax fails when it encounters an all-NaN slice in a chunk
Follow up on #776

```
In [1]: import numpy as np

In [2]: import dask.array as da

In [3]: x = np.array([[1.0, np.nan], [np.nan, 2.0]])

In [4]: da.nanmax(da.from_array(x, chunks=1), axis=1).compute()
/Users/shoyer/miniconda/envs/dask-dev/lib/python2.7/site-packages/numpy/lib/nanfunctions.py:319: RuntimeWarning: All-NaN slice encountered
  warnings.warn("All-NaN slice encountered", RuntimeWarning)
Out[4]: array([ 1.,  2.])

In [5]: da.nanargmax(da.from_array(x, chunks=1), axis=1).compute()
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-403b812e207c> in <module>()
----> 1 da.nanargmax(da.from_array(x, chunks=1), axis=1).compute()

/Users/shoyer/dev/dask/dask/base.pyc in compute(self, **kwargs)
     29
     30     def compute(self, **kwargs):
---> 31         return compute(self, **kwargs)[0]
     32
     33     @classmethod

/Users/shoyer/dev/dask/dask/base.pyc in compute(*args, **kwargs)
     97                 for opt, val in groups.items()])
     98     keys = [arg._keys() for arg in args]
---> 99     results = get(dsk, keys, **kwargs)
    100     return tuple(a._finalize(a, r) for a, r in zip(args, results))
    101

/Users/shoyer/dev/dask/dask/threaded.pyc in get(dsk, result, cache, num_workers, **kwargs)
     55     results = get_async(pool.apply_async, len(pool._pool), dsk, result,
     56                         cache=cache, queue=queue, get_id=_thread_get_id,
---> 57                         **kwargs)
     58
     59     return results

/Users/shoyer/dev/dask/dask/async.pyc in get_async(apply_async, num_workers, dsk, result, cache, queue, get_id, raise_on_exception, rerun_exceptions_locally, callbacks, **kwargs)
    480                 _execute_task(task, data)  # Re-execute locally
    481             else:
--> 482                 raise(remote_exception(res, tb))
    483         state['cache'][key] = res
    484         finish_task(dsk, key, state, results, keyorder.get)

ValueError: All-NaN slice encountered

Traceback
---------
  File "dask/async.py", line 262, in execute_task
    result = _execute_task(task, data)
  File "dask/async.py", line 245, in _execute_task
    return func(*args2)
  File "dask/array/reductions.py", line 367, in argreduce
    return (func(x, axis=axis), argfunc(x, axis=axis))
  File "dask/array/chunk.py", line 25, in keepdims_wrapped_callable
    r = a_callable(x, axis=axis, *args, **kwargs)
  File "/Users/shoyer/miniconda/envs/dask-dev/lib/python2.7/site-packages/numpy/lib/nanfunctions.py", line 420, in nanargmax
    raise ValueError("All-NaN slice encountered")

```

Discussion:

Patch:
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index a29003285..c0b12cd08 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -466,6 +466,13 @@ def arg_agg(func, argfunc, data, axis=None, **kwargs):
     return _arg_combine(data, axis, argfunc, keepdims=False)[0]
 
 
+def nanarg_agg(func, argfunc, data, axis=None, **kwargs):
+    arg, vals = _arg_combine(data, axis, argfunc, keepdims=False)
+    if np.any(np.isnan(vals)):
+        raise ValueError("All NaN slice encountered")
+    return arg
+
+
 def arg_reduction(x, chunk, combine, agg, axis=None, split_every=None):
     """Generic function for argreduction.
 
@@ -514,7 +521,7 @@ def arg_reduction(x, chunk, combine, agg, axis=None, split_every=None):
     return _tree_reduce(tmp, agg, axis, False, np.int64, split_every, combine)
 
 
-def make_arg_reduction(func, argfunc):
+def make_arg_reduction(func, argfunc, is_nan_func=False):
     """Create a argreduction callable.
 
     Parameters
@@ -526,17 +533,34 @@ def make_arg_reduction(func, argfunc):
     """
     chunk = partial(arg_chunk, func, argfunc)
     combine = partial(arg_combine, func, argfunc)
-    agg = partial(arg_agg, func, argfunc)
+    if is_nan_func:
+        agg = partial(nanarg_agg, func, argfunc)
+    else:
+        agg = partial(arg_agg, func, argfunc)
     @wraps(argfunc)
     def _(x, axis=None, split_every=None):
         return arg_reduction(x, chunk, combine, agg, axis, split_every)
     return _
 
 
+def _nanargmin(x, axis, **kwargs):
+    try:
+        return chunk.nanargmin(x, axis, **kwargs)
+    except ValueError:
+        return chunk.nanargmin(np.where(np.isnan(x), np.inf, x), axis, **kwargs)
+
+
+def _nanargmax(x, axis, **kwargs):
+    try:
+        return chunk.nanargmax(x, axis, **kwargs)
+    except ValueError:
+        return chunk.nanargmax(np.where(np.isnan(x), -np.inf, x), axis, **kwargs)
+
+
 argmin = make_arg_reduction(chunk.min, chunk.argmin)
 argmax = make_arg_reduction(chunk.max, chunk.argmax)
-nanargmin = make_arg_reduction(chunk.nanmin, chunk.nanargmin)
-nanargmax = make_arg_reduction(chunk.nanmax, chunk.nanargmax)
+nanargmin = make_arg_reduction(chunk.nanmin, _nanargmin, True)
+nanargmax = make_arg_reduction(chunk.nanmax, _nanargmax, True)
 
 
 def cumreduction(func, binop, ident, x, axis, dtype=None):


Test patch:
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index 7b734416f..2b1a08437 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -162,6 +162,26 @@ def test_arg_reductions(dfunc, func):
     assert eq(dfunc(a2, 0, split_every=2), func(x2, 0))
 
 
+@pytest.mark.parametrize(['dfunc', 'func'],
+         [(da.nanargmin, np.nanargmin), (da.nanargmax, np.nanargmax)])
+def test_nanarg_reductions(dfunc, func):
+    x = np.random.random((10, 10, 10))
+    x[5] = np.nan
+    a = da.from_array(x, chunks=(3, 4, 5))
+    assert eq(dfunc(a), func(x))
+    assert eq(dfunc(a, 0), func(x, 0))
+    with pytest.raises(ValueError):
+        dfunc(a, 1).compute()
+
+    with pytest.raises(ValueError):
+        dfunc(a, 2).compute()
+
+    x[:] = np.nan
+    a = da.from_array(x, chunks=(3, 4, 5))
+    with pytest.raises(ValueError):
+        dfunc(a).compute()
+
+
 def test_reductions_2D_nans():
     # chunks are a mix of some/all/no NaNs
     x = np.full((4, 4), np.nan)
@@ -189,17 +209,18 @@ def test_reductions_2D_nans():
     reduction_2d_test(da.nanmin, a, np.nanmin, x, False, False)
     reduction_2d_test(da.nanmax, a, np.nanmax, x, False, 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))
+    assert eq(da.argmax(a), np.argmax(x))
+    assert eq(da.argmin(a), np.argmin(x))
+    assert eq(da.nanargmax(a), np.nanargmax(x))
+    assert eq(da.nanargmin(a), np.nanargmin(x))
+    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():