Repository: cupy/cupy

Issue:
Add nanmean() nanstd()
It would be great if you can add functions like nanmean(), nanstd() and nanvar() like numpy


Discussion:
I would like to work on this issue.

Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -634,6 +634,7 @@ def isscalar(num):
 from cupy.statistics.meanvar import mean  # NOQA
 from cupy.statistics.meanvar import std  # NOQA
 from cupy.statistics.meanvar import var  # NOQA
+from cupy.statistics.meanvar import nanmean  # NOQA
 from cupy.statistics.meanvar import nanstd  # NOQA
 from cupy.statistics.meanvar import nanvar  # NOQA
 
diff --git a/cupy/statistics/meanvar.py b/cupy/statistics/meanvar.py
--- a/cupy/statistics/meanvar.py
+++ b/cupy/statistics/meanvar.py
@@ -139,7 +139,30 @@ def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
                  keepdims=keepdims)
 
 
-# TODO(okuta): Implement nanmean
+def nanmean(a, axis=None, dtype=None, out=None, keepdims=False):
+    """Returns the arithmetic mean along an axis ignoring NaN values.
+
+    Args:
+        a (cupy.ndarray): Array to compute mean.
+        axis (int): Along which axis to compute mean. The flattened array is
+            used by default.
+        dtype: Data type specifier.
+        out (cupy.ndarray): Output array.
+        keepdims (bool): If ``True``, the axis is remained as an axis of
+            size one.
+
+    Returns:
+        cupy.ndarray: The mean of the input array along the axis ignoring NaNs.
+
+    .. seealso:: :func:`numpy.nanmean`
+
+    """
+    if a.dtype.kind in 'biu':
+        return a.mean(axis=axis, dtype=dtype, out=out, keepdims=keepdims)
+
+    # TODO(okuta): check type
+    return _statistics._ndarray_nanmean(
+        a, axis=axis, dtype=dtype, out=out, keepdims=keepdims)
 
 
 def nanvar(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):


Test patch:
diff --git a/tests/cupy_tests/statics_tests/test_meanvar.py b/tests/cupy_tests/statics_tests/test_meanvar.py
--- a/tests/cupy_tests/statics_tests/test_meanvar.py
+++ b/tests/cupy_tests/statics_tests/test_meanvar.py
@@ -1,10 +1,14 @@
 import unittest
+import pytest
 
 import numpy
 
 import cupy
 from cupy import testing
 
+ignore_runtime_warnings = pytest.mark.filterwarnings(
+    "ignore", category=RuntimeWarning)
+
 
 @testing.gpu
 class TestAverage(unittest.TestCase):
@@ -181,6 +185,77 @@ def test_external_std_axis_ddof(self, xp, dtype):
         return xp.std(a, axis=1, ddof=1)
 
 
+@testing.parameterize(
+    *testing.product({
+        'shape': [(3, 4), (30, 40, 50)],
+        'axis': [None, 0, 1],
+        'keepdims': [True, False]
+    })
+)
+@testing.gpu
+class TestNanMean(unittest.TestCase):
+
+    @testing.for_all_dtypes(no_float16=True)
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_nanmean_without_nan(self, xp, dtype):
+        a = testing.shaped_random(self.shape, xp, dtype)
+        return xp.nanmean(a, axis=self.axis, keepdims=self.keepdims)
+
+    @ignore_runtime_warnings
+    @testing.for_all_dtypes(no_float16=True)
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_nanmean_with_nan_float(self, xp, dtype):
+        a = testing.shaped_random(self.shape, xp, dtype)
+
+        if a.dtype.kind not in 'biu':
+            a[1, :] = xp.nan
+            a[:, 3] = xp.nan
+
+        return xp.nanmean(a, axis=self.axis, keepdims=self.keepdims)
+
+
+@testing.gpu
+class TestNanMeanAdditional(unittest.TestCase):
+
+    @ignore_runtime_warnings
+    @testing.for_all_dtypes(no_float16=True)
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_nanmean_out(self, xp, dtype):
+        a = testing.shaped_random((10, 20, 30), xp, dtype)
+        z = xp.zeros((20, 30), dtype=dtype)
+
+        if a.dtype.kind not in 'biu':
+            a[1, :] = xp.nan
+            a[:, 3] = xp.nan
+
+        xp.nanmean(a, axis=0, out=z)
+        return z
+
+    @testing.slow
+    @testing.for_all_dtypes(no_float16=True)
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_nanmean_huge(self, xp, dtype):
+        a = testing.shaped_random((1024, 512), xp, dtype)
+
+        if a.dtype.kind not in 'biu':
+            a[:512, :256] = xp.nan
+
+        return xp.nanmean(a, axis=1)
+
+    @testing.numpy_cupy_allclose(rtol=1e-4)
+    def test_nanmean_float16(self, xp):
+        a = testing.shaped_arange((2, 3), xp, numpy.float16)
+        a[0][0] = xp.nan
+        return xp.nanmean(a)
+
+    @ignore_runtime_warnings
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_nanmean_all_nan(self, xp):
+        a = xp.zeros((3, 4))
+        a[:] = xp.nan
+        return xp.nanmean(a)
+
+
 @testing.parameterize(
     *testing.product({
         'shape': [(3, 4), (4, 3, 5)],
@@ -191,6 +266,7 @@ def test_external_std_axis_ddof(self, xp, dtype):
 @testing.gpu
 class TestNanVarStd(unittest.TestCase):
 
+    @ignore_runtime_warnings
     @testing.for_all_dtypes(no_float16=True, no_complex=True)
     @testing.numpy_cupy_allclose(rtol=1e-6)
     def test_nanvar(self, xp, dtype):
@@ -200,6 +276,7 @@ def test_nanvar(self, xp, dtype):
         return xp.nanvar(
             a, axis=self.axis, ddof=self.ddof, keepdims=self.keepdims)
 
+    @ignore_runtime_warnings
     @testing.for_all_dtypes(no_float16=True, no_complex=True)
     @testing.numpy_cupy_allclose(rtol=1e-6)
     def test_nanstd(self, xp, dtype):
@@ -213,6 +290,7 @@ def test_nanstd(self, xp, dtype):
 @testing.gpu
 class TestNanVarStdAdditional(unittest.TestCase):
 
+    @ignore_runtime_warnings
     @testing.for_all_dtypes(no_float16=True, no_complex=True)
     @testing.numpy_cupy_allclose(rtol=1e-6)
     def test_nanvar_out(self, xp, dtype):
@@ -243,6 +321,7 @@ def test_nanvar_float16(self, xp):
         a[0][0] = xp.nan
         return xp.nanvar(a, axis=0)
 
+    @ignore_runtime_warnings
     @testing.for_all_dtypes(no_float16=True, no_complex=True)
     @testing.numpy_cupy_allclose(rtol=1e-6)
     def test_nanstd_out(self, xp, dtype):