Repository: dask/dask

Issue:
Some nan functions don't exist for certain versions of NumPy
The following functions `nanmean`, `nanstd`, `nanvar` are not defined until NumPy 1.8.0. It would be nice if they weren't explicitly required.



Discussion:

Resolving patch:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -1,8 +1,87 @@
 """ A set of NumPy functions to apply per chunk """
 
+from __future__ import absolute_import, division, print_function
+
+from collections import Container, Iterable, Sequence
+from functools import wraps
+from inspect import getargspec
+
 from toolz import concat
 import numpy as np
 
+from ..utils import ignoring
+
+
+def keepdims_wrapper(a_callable):
+    """
+    A wrapper for functions that don't provide keepdims to ensure that they do.
+    """
+
+    if "keepdims" in getargspec(a_callable).args:
+        return a_callable
+
+    @wraps(a_callable)
+    def keepdims_wrapped_callable(x, axis=None, keepdims=None, *args, **kwargs):
+        r = a_callable(x, axis=axis, *args, **kwargs)
+
+        if not keepdims:
+            return r
+
+        axes = axis
+
+        if axes is None:
+            axes = range(x.ndim)
+
+        if not isinstance(axes, (Container, Iterable, Sequence)):
+            axes = [axes]
+
+        r_slice = tuple()
+        for each_axis in range(x.ndim):
+            if each_axis in axes:
+                r_slice += (None,)
+            else:
+                r_slice += (slice(None),)
+
+        r = r[r_slice]
+
+        return r
+
+    return keepdims_wrapped_callable
+
+
+# Wrap NumPy functions to ensure they provide keepdims.
+sum = keepdims_wrapper(np.sum)
+prod = keepdims_wrapper(np.prod)
+min = keepdims_wrapper(np.min)
+max = keepdims_wrapper(np.max)
+argmin = keepdims_wrapper(np.argmin)
+nanargmin = keepdims_wrapper(np.nanargmin)
+argmax = keepdims_wrapper(np.argmax)
+nanargmax = keepdims_wrapper(np.nanargmax)
+any = keepdims_wrapper(np.any)
+all = keepdims_wrapper(np.all)
+nansum = keepdims_wrapper(np.nansum)
+
+with ignoring(AttributeError):
+    nanprod = keepdims_wrapper(np.nanprod)
+
+nanmin = keepdims_wrapper(np.nanmin)
+nanmax = keepdims_wrapper(np.nanmax)
+mean = keepdims_wrapper(np.mean)
+
+with ignoring(AttributeError):
+    nanmean = keepdims_wrapper(np.nanmean)
+
+var = keepdims_wrapper(np.var)
+
+with ignoring(AttributeError):
+    nanvar = keepdims_wrapper(np.nanvar)
+
+std = keepdims_wrapper(np.std)
+
+with ignoring(AttributeError):
+    nanstd = keepdims_wrapper(np.nanstd)
+
 
 def coarsen(reduction, x, axes):
     """ Coarsen array by applying reduction to fixed size neighborhoods
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -7,6 +7,7 @@
 from .core import (_concatenate2, insert_many, Array, atop, names, sqrt,
         elemwise)
 from ..core import flatten
+from . import chunk
 from ..utils import ignoring
 
 
@@ -42,88 +43,88 @@ def reduction(x, chunk, aggregate, axis=None, keepdims=None):
         return result
 
 
-@wraps(np.sum)
+@wraps(chunk.sum)
 def sum(a, axis=None, keepdims=False):
-    return reduction(a, np.sum, np.sum, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.sum, chunk.sum, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.prod)
+@wraps(chunk.prod)
 def prod(a, axis=None, keepdims=False):
-    return reduction(a, np.prod, np.prod, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.prod, chunk.prod, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.min)
+@wraps(chunk.min)
 def min(a, axis=None, keepdims=False):
-    return reduction(a, np.min, np.min, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.min, chunk.min, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.max)
+@wraps(chunk.max)
 def max(a, axis=None, keepdims=False):
-    return reduction(a, np.max, np.max, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.max, chunk.max, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.argmin)
+@wraps(chunk.argmin)
 def argmin(a, axis=None):
-    return arg_reduction(a, np.min, np.argmin, axis=axis)
+    return arg_reduction(a, chunk.min, chunk.argmin, axis=axis)
 
 
-@wraps(np.nanargmin)
+@wraps(chunk.nanargmin)
 def nanargmin(a, axis=None):
-    return arg_reduction(a, np.nanmin, np.nanargmin, axis=axis)
+    return arg_reduction(a, chunk.nanmin, chunk.nanargmin, axis=axis)
 
 
-@wraps(np.argmax)
+@wraps(chunk.argmax)
 def argmax(a, axis=None):
-    return arg_reduction(a, np.max, np.argmax, axis=axis)
+    return arg_reduction(a, chunk.max, chunk.argmax, axis=axis)
 
 
-@wraps(np.nanargmax)
+@wraps(chunk.nanargmax)
 def nanargmax(a, axis=None):
-    return arg_reduction(a, np.nanmax, np.nanargmax, axis=axis)
+    return arg_reduction(a, chunk.nanmax, chunk.nanargmax, axis=axis)
 
 
-@wraps(np.any)
+@wraps(chunk.any)
 def any(a, axis=None, keepdims=False):
-    return reduction(a, np.any, np.any, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.any, chunk.any, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.all)
+@wraps(chunk.all)
 def all(a, axis=None, keepdims=False):
-    return reduction(a, np.all, np.all, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.all, chunk.all, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.nansum)
+@wraps(chunk.nansum)
 def nansum(a, axis=None, keepdims=False):
-    return reduction(a, np.nansum, np.sum, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.nansum, chunk.sum, axis=axis, keepdims=keepdims)
 
 
 with ignoring(AttributeError):
-    @wraps(np.nanprod)
+    @wraps(chunk.nanprod)
     def nanprod(a, axis=None, keepdims=False):
-        return reduction(a, np.nanprod, np.prod, axis=axis, keepdims=keepdims)
+        return reduction(a, chunk.nanprod, chunk.prod, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.nanmin)
+@wraps(chunk.nanmin)
 def nanmin(a, axis=None, keepdims=False):
-    return reduction(a, np.nanmin, np.min, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.nanmin, chunk.min, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.nanmax)
+@wraps(chunk.nanmax)
 def nanmax(a, axis=None, keepdims=False):
-    return reduction(a, np.nanmax, np.max, axis=axis, keepdims=keepdims)
+    return reduction(a, chunk.nanmax, chunk.max, axis=axis, keepdims=keepdims)
 
 
 def numel(x, **kwargs):
     """ A reduction to count the number of elements """
-    return np.sum(np.ones_like(x), **kwargs)
+    return chunk.sum(np.ones_like(x), **kwargs)
 
 
 def nannumel(x, **kwargs):
     """ A reduction to count the number of elements """
-    return np.sum(~np.isnan(x), **kwargs)
+    return chunk.sum(~np.isnan(x), **kwargs)
 
 
-def mean_chunk(x, sum=np.sum, numel=numel, **kwargs):
+def mean_chunk(x, sum=chunk.sum, numel=numel, **kwargs):
     n = numel(x, **kwargs)
     total = sum(x, **kwargs)
     result = np.empty(shape=n.shape,
@@ -136,17 +137,18 @@ def mean_agg(pair, **kwargs):
     return pair['total'].sum(**kwargs) / pair['n'].sum(**kwargs)
 
 
-@wraps(np.mean)
+@wraps(chunk.mean)
 def mean(a, axis=None, keepdims=False):
     return reduction(a, mean_chunk, mean_agg, axis=axis, keepdims=keepdims)
 
 
-@wraps(np.nanmean)
 def nanmean(a, axis=None, keepdims=False):
-    return reduction(a, partial(mean_chunk, sum=np.nansum, numel=nannumel),
+    return reduction(a, partial(mean_chunk, sum=chunk.nansum, numel=nannumel),
                      mean_agg, axis=axis, keepdims=keepdims)
+with ignoring(AttributeError):
+    nanmean = wraps(chunk.nanmean)(nanmean)
 
-def var_chunk(A, sum=np.sum, numel=numel, **kwargs):
+def var_chunk(A, sum=chunk.sum, numel=numel, **kwargs):
     n = numel(A, **kwargs)
     x = sum(A, dtype='f8', **kwargs)
     x2 = sum(A**2, dtype='f8', **kwargs)
@@ -168,24 +170,26 @@ def var_agg(A, ddof=None, **kwargs):
     return result
 
 
-@wraps(np.var)
+@wraps(chunk.var)
 def var(a, axis=None, keepdims=False, ddof=0):
     return reduction(a, var_chunk, partial(var_agg, ddof=ddof), axis=axis, keepdims=keepdims)
 
 
-@wraps(np.nanvar)
 def nanvar(a, axis=None, keepdims=False, ddof=0):
-    return reduction(a, partial(var_chunk, sum=np.nansum, numel=nannumel),
+    return reduction(a, partial(var_chunk, sum=chunk.nansum, numel=nannumel),
                      partial(var_agg, ddof=ddof), axis=axis, keepdims=keepdims)
+with ignoring(AttributeError):
+    nanvar = wraps(chunk.nanvar)(nanvar)
 
-@wraps(np.std)
+@wraps(chunk.std)
 def std(a, axis=None, keepdims=False, ddof=0):
     return sqrt(a.var(axis=axis, keepdims=keepdims, ddof=ddof))
 
 
-@wraps(np.nanstd)
 def nanstd(a, axis=None, keepdims=False, ddof=0):
     return sqrt(nanvar(a, axis=axis, keepdims=keepdims, ddof=ddof))
+with ignoring(AttributeError):
+    nanstd = wraps(chunk.nanstd)(nanstd)
 
 
 def vnorm(a, ord=None, axis=None, keepdims=False):


Test patch:
diff --git a/dask/array/tests/test_chunk.py b/dask/array/tests/test_chunk.py
--- a/dask/array/tests/test_chunk.py
+++ b/dask/array/tests/test_chunk.py
@@ -1,5 +1,93 @@
+from __future__ import absolute_import, division, print_function
+
 import numpy as np
-from dask.array.chunk import coarsen, trim
+
+from dask.array.chunk import coarsen, keepdims_wrapper, trim
+
+
+def test_keepdims_wrapper_no_axis():
+    def summer(a, axis=None):
+        return a.sum(axis=axis)
+
+    summer_wrapped = keepdims_wrapper(summer)
+
+    assert summer_wrapped != summer
+    assert summer_wrapped == keepdims_wrapper(summer_wrapped)
+
+    a = np.arange(24).reshape(1, 2, 3, 4)
+
+    r = summer(a)
+    rw = summer_wrapped(a, keepdims=True)
+    rwf = summer_wrapped(a, keepdims=False)
+
+    assert r.ndim == 0
+    assert r.shape == tuple()
+    assert r == 276
+
+    assert rw.ndim == 4
+    assert rw.shape == (1, 1, 1, 1)
+    assert (rw == 276).all()
+
+    assert rwf.ndim == 0
+    assert rwf.shape == tuple()
+    assert rwf == 276
+
+
+def test_keepdims_wrapper_one_axis():
+    def summer(a, axis=None):
+        return a.sum(axis=axis)
+
+    summer_wrapped = keepdims_wrapper(summer)
+
+    assert summer_wrapped != summer
+    assert summer_wrapped == keepdims_wrapper(summer_wrapped)
+
+    a = np.arange(24).reshape(1, 2, 3, 4)
+
+    r = summer(a, axis=2)
+    rw = summer_wrapped(a, axis=2, keepdims=True)
+    rwf = summer_wrapped(a, axis=2, keepdims=False)
+
+    assert r.ndim == 3
+    assert r.shape == (1, 2, 4)
+    assert (r == np.array([[[12, 15, 18, 21], [48, 51, 54, 57]]])).all()
+
+    assert rw.ndim == 4
+    assert rw.shape == (1, 2, 1, 4)
+    assert (rw == np.array([[[[12, 15, 18, 21]], [[48, 51, 54, 57]]]])).all()
+
+    assert rwf.ndim == 3
+    assert rwf.shape == (1, 2, 4)
+    assert (rwf == np.array([[[12, 15, 18, 21], [48, 51, 54, 57]]])).all()
+
+
+def test_keepdims_wrapper_two_axes():
+    def summer(a, axis=None):
+        return a.sum(axis=axis)
+
+    summer_wrapped = keepdims_wrapper(summer)
+
+    assert summer_wrapped != summer
+    assert summer_wrapped == keepdims_wrapper(summer_wrapped)
+
+    a = np.arange(24).reshape(1, 2, 3, 4)
+
+    r = summer(a, axis=(1, 3))
+    rw = summer_wrapped(a, axis=(1, 3), keepdims=True)
+    rwf = summer_wrapped(a, axis=(1, 3), keepdims=False)
+
+    assert r.ndim == 2
+    assert r.shape == (1, 3)
+    assert (r == np.array([[60, 92, 124]])).all()
+
+    assert rw.ndim == 4
+    assert rw.shape == (1, 1, 3, 1)
+    assert (rw == np.array([[[[60], [92], [124]]]])).all()
+
+    assert rwf.ndim == 2
+    assert rwf.shape == (1, 3)
+    assert (rwf == np.array([[60, 92, 124]])).all()
+
 
 def eq(a, b):
     c = a == b