Repository: scipy/scipy

Issue:
NaN handling of stats.rankdata
I think NaN handling of stats.rankdata is not correct. 

#### Reproducing code example:

Current code is

```
rankdata([0, 2, 3, 2, np.nan, np.nan])
> [1.  2.5 4.  2.5 5.  6. ]
```

However, I think it should be

```
rankdata([0, 2, 3, 2, np.nan, np.nan])
> [1.  2.5 4.  2.5 np.nan  np.nan ]
```

or 

return nan

or

raise an exception.

#### Scipy/Numpy/Python version information:
Python: 3.7.7
scipy: 1.7.1


Discussion:
see also the related issue #11409
@chrisb83 Oh. Thanks. I will try to fix it in #11791

Resolving patch:
diff --git a/scipy/stats/_stats_py.py b/scipy/stats/_stats_py.py
--- a/scipy/stats/_stats_py.py
+++ b/scipy/stats/_stats_py.py
@@ -9030,7 +9030,7 @@ def _square_of_sums(a, axis=0):
         return float(s) * s
 
 
-def rankdata(a, method='average', *, axis=None):
+def rankdata(a, method='average', *, axis=None, nan_policy='omit'):
     """Assign ranks to data, dealing with ties appropriately.
 
     By default (``axis=None``), the data array is first flattened, and a flat
@@ -9063,6 +9063,25 @@ def rankdata(a, method='average', *, axis=None):
     axis : {None, int}, optional
         Axis along which to perform the ranking. If ``None``, the data array
         is first flattened.
+    nan_policy : {'omit', 'propagate', 'raise'}, optional
+        Defines how to handle when input contains nan.
+        The following options are available (default is 'omit'):
+
+          * 'omit': performs the calculations ignoring nan values
+          * 'propagate': propagates nans through the rank calculation
+          * 'raise': raises an error
+
+        .. note::
+
+            When `nan_policy` is 'omit', nans in `a` are ignored when ranking
+            the other values, and the corresponding locations of the output
+            are nan. This behavior is the default because it is intuitive and
+            compatible with the behavior before the `nan_policy` parameter
+            was introduced.
+            When `nan_policy` is 'propagate', the output is an array of *all*
+            nans because ranks relative to nans in the input are undefined.
+
+        .. versionadded:: 1.10
 
     Returns
     -------
@@ -9093,13 +9112,18 @@ def rankdata(a, method='average', *, axis=None):
     >>> rankdata([[0, 2, 2], [3, 2, 5]], axis=1)
     array([[1. , 2.5, 2.5],
            [2. , 1. , 3. ]])
+    >>> rankdata([0, 2, 3, np.nan, -2, np.nan], nan_policy="omit")
+    array([ 2.,  3.,  4., nan,  1., nan])
+    >>> rankdata([0, 2, 3, np.nan, -2, np.nan], nan_policy="propagate")
+    array([nan, nan, nan, nan, nan, nan])
 
     """
     if method not in ('average', 'min', 'max', 'dense', 'ordinal'):
         raise ValueError('unknown method "{0}"'.format(method))
 
+    a = np.asarray(a)
+
     if axis is not None:
-        a = np.asarray(a)
         if a.size == 0:
             # The return values of `normalize_axis_index` are ignored.  The
             # call validates `axis`, even though we won't use it.
@@ -9107,9 +9131,18 @@ def rankdata(a, method='average', *, axis=None):
             np.core.multiarray.normalize_axis_index(axis, a.ndim)
             dt = np.float64 if method == 'average' else np.int_
             return np.empty(a.shape, dtype=dt)
-        return np.apply_along_axis(rankdata, axis, a, method)
+        return np.apply_along_axis(rankdata, axis, a, method,
+                                   nan_policy=nan_policy)
+
+    contains_nan, nan_policy = _contains_nan(a, nan_policy)
+    nan_indexes = None
+    if contains_nan:
+        if nan_policy == 'omit':
+            nan_indexes = np.isnan(a)
+        if nan_policy == 'propagate':
+            return np.full_like(a, np.nan)
 
-    arr = np.ravel(np.asarray(a))
+    arr = np.ravel(a)
     algo = 'mergesort' if method == 'ordinal' else 'quicksort'
     sorter = np.argsort(arr, kind=algo)
 
@@ -9117,23 +9150,29 @@ def rankdata(a, method='average', *, axis=None):
     inv[sorter] = np.arange(sorter.size, dtype=np.intp)
 
     if method == 'ordinal':
-        return inv + 1
+        result = inv + 1
 
     arr = arr[sorter]
     obs = np.r_[True, arr[1:] != arr[:-1]]
     dense = obs.cumsum()[inv]
 
     if method == 'dense':
-        return dense
+        result = dense
 
     # cumulative counts of each unique value
     count = np.r_[np.nonzero(obs)[0], len(obs)]
 
     if method == 'max':
-        return count[dense]
+        result = count[dense]
 
     if method == 'min':
-        return count[dense - 1] + 1
+        result = count[dense - 1] + 1
+
+    if method == 'average':
+        result = .5 * (count[dense] + count[dense - 1] + 1)
+
+    if nan_indexes is not None:
+        result = result.astype('float64')
+        result[nan_indexes] = np.nan
 
-    # average method
-    return .5 * (count[dense] + count[dense - 1] + 1)
+    return result


Test patch:
diff --git a/scipy/stats/tests/test_rank.py b/scipy/stats/tests/test_rank.py
--- a/scipy/stats/tests/test_rank.py
+++ b/scipy/stats/tests/test_rank.py
@@ -192,6 +192,74 @@ def test_size_0_axis(self, axis, method, dtype):
         assert_equal(r.shape, shape)
         assert_equal(r.dtype, dtype)
 
+    @pytest.mark.parametrize('axis', range(3))
+    @pytest.mark.parametrize('method', methods)
+    def test_nan_policy_omit_3d(self, axis, method):
+        shape = (20, 21, 22)
+        rng = np.random.default_rng(abs(hash('falafel')))
+
+        a = rng.random(size=shape)
+        i = rng.random(size=shape) < 0.4
+        j = rng.random(size=shape) < 0.1
+        k = rng.random(size=shape) < 0.1
+        a[i] = np.nan
+        a[j] = -np.inf
+        a[k] - np.inf
+
+        def rank_1d_omit(a, method):
+            out = np.zeros_like(a)
+            i = np.isnan(a)
+            a_compressed = a[~i]
+            res = rankdata(a_compressed, method)
+            out[~i] = res
+            out[i] = np.nan
+            return out
+
+        def rank_omit(a, method, axis):
+            return np.apply_along_axis(lambda a: rank_1d_omit(a, method),
+                                       axis, a)
+
+        res = rankdata(a, method, axis=axis, nan_policy='omit')
+        res0 = rank_omit(a, method, axis=axis)
+
+        assert_array_equal(res, res0)
+
+    def test_nan_policy_raise(self):
+        # 1 1d-array test
+        data = [0, 2, 3, -2, np.nan, np.nan]
+        with pytest.raises(ValueError, match="The input contains nan"):
+            rankdata(data, nan_policy='raise')
+
+        # 2 2d-array test
+        data = [[0, np.nan, 3],
+                [4, 2, np.nan],
+                [np.nan, 2, 2]]
+
+        with pytest.raises(ValueError, match="The input contains nan"):
+            rankdata(data, axis=0, nan_policy="raise")
+
+        with pytest.raises(ValueError, match="The input contains nan"):
+            rankdata(data, axis=1, nan_policy="raise")
+
+    def test_nan_policy_propagate(self):
+        # 1 1d-array test
+        data = [0, 2, 3, -2, np.nan, np.nan]
+        assert_array_equal(rankdata(data, nan_policy='propagate'),
+                           [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan])
+
+        # 2 2d-array test
+        data = [[0, np.nan, 3],
+                [4, 2, np.nan],
+                [1, 2, 2]]
+        assert_array_equal(rankdata(data, axis=0, nan_policy='propagate'),
+                           [[1, np.nan, np.nan],
+                            [3, np.nan, np.nan],
+                            [2, np.nan, np.nan]])
+        assert_array_equal(rankdata(data, axis=1, nan_policy='propagate'),
+                           [[np.nan, np.nan, np.nan],
+                            [np.nan, np.nan, np.nan],
+                            [1, 2.5, 2.5]])
+
 
 _cases = (
     # values, method, expected