TST: Added tests for nanmean(), nanvar(), nanstd()

2 files changed, 134 insertions, 10 deletions

diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py
index 297d708a8..c317358e1 100644
--- a/numpy/core/_methods.py
+++ b/numpy/core/_methods.py
@@ -7,7 +7,8 @@ from __future__ import division, absolute_import, print_function
 
 from numpy.core import multiarray as mu
 from numpy.core import umath as um
-from numpy.core.numeric import array, asanyarray, isnan
+from numpy.core.numeric import array, asanyarray, isnan, issubdtype
+from numpy.core import numerictypes as nt
 
 def _amax(a, axis=None, out=None, keepdims=False):
     return um.maximum.reduce(a, axis=axis,
@@ -46,8 +47,9 @@ def _count_reduce_items(arr, axis):
 def _mean(a, axis=None, dtype=None, out=None, keepdims=False):
     arr = asanyarray(a)
 
-    # Upgrade bool, unsigned int, and int to float64
-    if dtype is None and arr.dtype.kind in ['b','u','i']:
+    # Cast bool, unsigned int, and int to float64
+    if dtype is None and (issubdtype(arr.dtype, nt.integer) or
+                          issubdtype(arr.dtype, nt.bool_)):
         ret = um.add.reduce(arr, axis=axis, dtype='f8',
                             out=out, keepdims=keepdims)
     else:
@@ -62,11 +64,14 @@ def _mean(a, axis=None, dtype=None, out=None, keepdims=False):
     return ret
 
 def _nanmean(a, axis=None, dtype=None, out=None, keepdims=False):
+    # Using array() instead of asanyarray() because the former always
+    # makes a copy, which is important due to the copyto() action later
     arr = array(a, subok=True)
     mask = isnan(arr)
 
-    # Upgrade bool, unsigned int, and int to float64
-    if dtype is None and arr.dtype.kind in ['b','u','i']:
+    # Cast bool, unsigned int, and int to float64
+    if dtype is None and (issubdtype(arr.dtype, nt.integer) or
+                          issubdtype(arr.dtype, nt.bool_)):
         ret = um.add.reduce(arr, axis=axis, dtype='f8',
                             out=out, keepdims=keepdims)
     else:
@@ -86,7 +91,8 @@ def _var(a, axis=None, dtype=None, out=None, ddof=0,
     arr = asanyarray(a)
 
     # First compute the mean, saving 'rcount' for reuse later
-    if dtype is None and arr.dtype.kind in ['b','u','i']:
+    if dtype is None and (issubdtype(arr.dtype, nt.integer) or
+                          issubdtype(arr.dtype, nt.bool_)):
         arrmean = um.add.reduce(arr, axis=axis, dtype='f8', keepdims=True)
     else:
         arrmean = um.add.reduce(arr, axis=axis, dtype=dtype, keepdims=True)
@@ -101,7 +107,7 @@ def _var(a, axis=None, dtype=None, out=None, ddof=0,
     x = arr - arrmean
 
     # (arr - arrmean) ** 2
-    if arr.dtype.kind == 'c':
+    if issubdtype(arr.dtype, nt.complex_):
         x = um.multiply(x, um.conjugate(x), out=x).real
     else:
         x = um.multiply(x, x, out=x)
@@ -123,11 +129,14 @@ def _var(a, axis=None, dtype=None, out=None, ddof=0,
 
 def _nanvar(a, axis=None, dtype=None, out=None, ddof=0,
                             keepdims=False):
+    # Using array() instead of asanyarray() because the former always
+    # makes a copy, which is important due to the copyto() action later
     arr = array(a, subok=True)
     mask = isnan(arr)
 
     # First compute the mean, saving 'rcount' for reuse later
-    if dtype is None and arr.dtype.kind in ['b','u','i']:
+    if dtype is None and (issubdtype(arr.dtype, nt.integer) or
+                          issubdtype(arr.dtype, nt.bool_)):
         arrmean = um.add.reduce(arr, axis=axis, dtype='f8', keepdims=True)
     else:
         mu.copyto(arr, 0.0, where=mask)
@@ -142,10 +151,10 @@ def _nanvar(a, axis=None, dtype=None, out=None, ddof=0,
 
     # arr - arrmean
     x = arr - arrmean
-    x[mask] = 0.0
+    mu.copyto(x, 0.0, where=mask)
 
     # (arr - arrmean) ** 2
-    if arr.dtype.kind == 'c':
+    if issubdtype(arr.dtype, nt.complex_):
         x = um.multiply(x, um.conjugate(x), out=x).real
     else:
         x = um.multiply(x, x, out=x)
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index 782ddd687..0d02ce6c6 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -3,6 +3,7 @@ from __future__ import division, absolute_import, print_function
 import sys
 import platform
 from decimal import Decimal
+import warnings
 
 import numpy as np
 from numpy.core import *
@@ -177,18 +178,51 @@ class TestNonarrayArgs(TestCase):
         assert_(all(mean(A,0) == array([2.5,3.5,4.5])))
         assert_(all(mean(A,1) == array([2.,5.])))
 
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            assert_(isnan(mean([])))
+            assert_(w[0].category is RuntimeWarning)
+
+    def test_nanmean(self):
+        A = [[1, nan, nan], [nan, 4, 5]]
+        assert_(nanmean(A) == (10.0 / 3))
+        assert_(all(nanmean(A,0) == array([1, 4, 5])))
+        assert_(all(nanmean(A,1) == array([1, 4.5])))
+
     def test_std(self):
         A = [[1,2,3],[4,5,6]]
         assert_almost_equal(std(A), 1.707825127659933)
         assert_almost_equal(std(A,0), array([1.5, 1.5, 1.5]))
         assert_almost_equal(std(A,1), array([0.81649658, 0.81649658]))
 
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            assert_(isnan(std([])))
+            assert_(w[0].category is RuntimeWarning)
+
+    def test_nanstd(self):
+        A = [[1, nan, nan], [nan, 4, 5]]
+        assert_almost_equal(nanstd(A), 1.699673171197595)
+        assert_almost_equal(nanstd(A,0), array([0.0, 0.0, 0.0]))
+        assert_almost_equal(nanstd(A,1), array([0.0, 0.5]))
+
     def test_var(self):
         A = [[1,2,3],[4,5,6]]
         assert_almost_equal(var(A), 2.9166666666666665)
         assert_almost_equal(var(A,0), array([2.25, 2.25, 2.25]))
         assert_almost_equal(var(A,1), array([0.66666667, 0.66666667]))
 
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            assert_(isnan(mean([])))
+            assert_(w[0].category is RuntimeWarning)
+
+    def test_nanvar(self):
+        A = [[1, nan, nan], [nan, 4, 5]]
+        assert_almost_equal(nanvar(A), 2.88888888889)
+        assert_almost_equal(nanvar(A,0), array([0.0, 0.0, 0.0]))
+        assert_almost_equal(nanvar(A,1), array([0.0, 0.25]))
+
 
 class TestBoolScalar(TestCase):
     def test_logical(self):
@@ -1447,6 +1481,39 @@ class TestIsclose(object):
         assert_array_equal(x, array([inf, 1]))
         assert_array_equal(y, array([0, inf]))
 
+class TestNaNMean(TestCase):
+    def setUp(self):
+        self.A = array([1,nan,-1,nan,nan,1,-1])
+        self.B = array([nan, nan, nan, nan])
+        self.real_mean = 0
+
+    def test_basic(self):
+        assert_almost_equal(nanmean(self.A),self.real_mean)
+
+    def test_mutation(self):
+        # Because of the "messing around" we do to replace NaNs with zeros
+        # this is meant to ensure we don't actually replace the NaNs in the
+        # actual array.
+        a_copy = self.A.copy()
+        b_copy = self.B.copy()
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            a_ret = nanmean(self.A)
+            assert_equal(self.A, a_copy)
+            b_ret = nanmean(self.B)
+            assert_equal(self.B, b_copy)
+
+    def test_allnans(self):
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            assert_(isnan(nanmean(self.B)))
+            assert_(w[0].category is RuntimeWarning)
+
+    def test_empty(self):
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            assert_(isnan(nanmean(array([]))))
+            assert_(w[0].category is RuntimeWarning)
 
 class TestStdVar(TestCase):
     def setUp(self):
@@ -1469,6 +1536,54 @@ class TestStdVar(TestCase):
         assert_almost_equal(std(self.A,ddof=2)**2,
                             self.real_var*len(self.A)/float(len(self.A)-2))
 
+class TestNaNStdVar(TestCase):
+    def setUp(self):
+        self.A = array([nan,1,-1,nan,1,nan,-1])
+        self.B = array([nan, nan, nan, nan])
+        self.real_var = 1
+
+    def test_basic(self):
+        assert_almost_equal(nanvar(self.A),self.real_var)
+        assert_almost_equal(nanstd(self.A)**2,self.real_var)
+
+    def test_mutation(self):
+        # Because of the "messing around" we do to replace NaNs with zeros
+        # this is meant to ensure we don't actually replace the NaNs in the
+        # actual array.
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            a_copy = self.A.copy()
+            b_copy = self.B.copy()
+            a_ret = nanvar(self.A)
+            assert_equal(self.A, a_copy)
+            b_ret = nanstd(self.B)
+            assert_equal(self.B, b_copy)
+
+    def test_ddof1(self):
+        assert_almost_equal(nanvar(self.A,ddof=1),
+                self.real_var*sum(~isnan(self.A))/float(sum(~isnan(self.A))-1))
+        assert_almost_equal(nanstd(self.A,ddof=1)**2,
+                self.real_var*sum(~isnan(self.A))/float(sum(~isnan(self.A))-1))
+
+    def test_ddof2(self):
+        assert_almost_equal(nanvar(self.A,ddof=2),
+                self.real_var*sum(~isnan(self.A))/float(sum(~isnan(self.A))-2))
+        assert_almost_equal(nanstd(self.A,ddof=2)**2,
+                self.real_var*sum(~isnan(self.A))/float(sum(~isnan(self.A))-2))
+
+    def test_allnans(self):
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            assert_(isnan(nanvar(self.B)))
+            assert_(isnan(nanstd(self.B)))
+            assert_(w[0].category is RuntimeWarning)
+
+    def test_empty(self):
+        with warnings.catch_warnings(record=True) as w:
+            warnings.filterwarnings('always', '', RuntimeWarning)
+            assert_(isnan(nanvar(array([]))))
+            assert_(isnan(nanstd(array([]))))
+            assert_(w[0].category is RuntimeWarning)
 
 class TestStdVarComplex(TestCase):
     def test_basic(self):