Repository: scipy/scipy
Pull request: ENH: stats.mannwhitneyu: vectorize statistic calculation
#### Reference issue
Closes gh-19642

#### What does this implement/fix?
gh-19642 asked to speed up `stats.mannwhitneyu`. The bottleneck there appeared to be computation of the statistic itself (whereas in gh-19692 the problem was with the exact p-value). This PR vectorizes the calculation of data ranks and computes the tie correction term at the same time.

#### Additional information
Speed comparison:

```python3
import numpy as np
from scipy import stats
rng = np.random.default_rng(2549824598234528)
x = rng.integers(5, size=10)
y = rng.integers(5, size=20000)
%timeit stats.mannwhitneyu(x, y)
# Before: 9.54 ms ± 12.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# After: 1.01 ms ± 3.26 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

x = rng.integers(5, size=10)
y = rng.integers(5, size=4000)
%timeit stats.mannwhitneyu(x, y)
# Before: 2.22 ms ± 602 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# After: 421 µs ± 5.49 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

x = rng.integers(5, size=(10, 1))
y = rng.integers(5, size=(4000, 20000))
%timeit -r 1 -n 1 stats.mannwhitneyu(x, y)
# Before: 34.3 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
# After: 8.25 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

x = rng.integers(5, size=(10, 1))
y = rng.integers(5, size=(40, 20000))
%timeit stats.mannwhitneyu(x, y)
# Before: 1.38 s ± 4.85 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# After: 64.4 ms ± 595 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

Problem statement:
Speeding up Mann-Whitney U-Test

The current Mann-Whitney U-Test is very fast, but it's still slow for large datasets.  As there are tens of CPU for some servers, is it possible to speed up using some high performance techinques such as multi-threading and paralleling  to use more CPUs on the server.

Hints from issue discussion:
> The current Mann-Whitney U-Test is very fast

Thanks for noticing!

> is it possible to speed up...?

There are a few things I know offhand, like there is at least a factor of two in the calculation of exact p-values, and the ranking calculation is not natively vectorized. But which answer is relevant depends on what your bottleneck is. What array shapes do you have, what is `axis`, are there any ties, and do you need exact p-values (or are approximate p-values OK)? 
Hi, this is the code I am using: 
 `U1, p_value = stats.mannwhitneyu(x, y, nan_policy='omit')`
x is a list with 10 elements, y is a pandas dataframe with the shape (4000x20000). The values of x are tested for each column of y. there are many ties. Both exact p-values and approximate p-values are OK for me.
You wrote  "The values of x are tested for each column of y", but that is not what this code would do without shaping `x` accordingly and using `axis=1`.

~~Are you looping over the columns manually (i.e. calling `mannwhitneyu` 20000 times)?~~ never mind, with `nan_policy='omit'` it's going to have an explicit loop whatever you do.

Do both `x` and `y` have NaNs? What are the `dtype`s? (I ask because pandas integers support NaNs, too.) Are there the same number of NaNs in each column?
Hi, I made an error description!

In my case, the x is a dataframe with only one column or the same number of columns as y.
here is an toy data I used:

```python
import pandas as pd
import numpy as np
import scipy.stats as stats

# x = pd.DataFrame({'c1': [1.0, 2, 3, 4, 5]})
x = pd.DataFrame({'c1': [1, 2, 3, 4, 5],'c2': [11, 12, 13, 14, 15],'c3': [1, 2, 0, 4, 5],'c4': [1, 2, 3, np.nan, np.nan],'c5': [1, 2, 3, 4, 5]})
y = pd.DataFrame({
    'sample1':[19, 22, np.nan, 29, 24, 25],
    'sample2':[1, 2, 3, 4, 5, 6],
    'sample3':[19, 0, 20, 11, 17.0, 12.32],
    'sample4':[1,1, 3, 24, np.nan, 25],
    'sample5':[5,6,7,8,9,20]},
    index=['gene1','gene2','gene3','gene4','gene5','gene6']
)

U1, p_value = stats.mannwhitneyu(x, y, nan_policy='omit')
```
Both x and y have NaNs and the `dtype` are float64. The number of NaNs in each column is not always the same.
When y gets larger (shape 4000 * 20000), the u-test will spend about 60 seconds.
If you're willing to edit your source code, you will probably get a substantial speedup by modifying these lines:

https://github.com/scipy/scipy/blob/5e4a5e3785f79dd4e8930eed883da89958860db2/scipy/stats/_mannwhitneyu.py#L173-L179

Just comment out the first tie correction part and uncomment the second one. Please verify that the results you get are the same either way. If so, and if you observe improved performance, I'll look into making the change. 

We could probably speed that up further by vectorizing it rather than relying on `np.apply_along_axis`, but since you have NaNs, it would require adding native support for omitting NaNs throughout the code (rather than relying on a generic routine that loops over 1D slices with NaNs removed).

Commit message: MAINT: stats.mannwhitneyu: use new _rankdata

Patch to review:
diff --git a/scipy/stats/_mannwhitneyu.py b/scipy/stats/_mannwhitneyu.py
index 342eccbacf87..abbc25ecda59 100644
--- a/scipy/stats/_mannwhitneyu.py
+++ b/scipy/stats/_mannwhitneyu.py
@@ -2,6 +2,7 @@
 from collections import namedtuple
 from scipy import special
 from scipy import stats
+from scipy.stats._stats_py import _rankdata
 from ._axis_nan_policy import _axis_nan_policy_factory
 
 
@@ -157,27 +158,17 @@ def _base_case(m, n, k):
     return fmnks
 
 
-def _tie_term(ranks):
-    """Tie correction term"""
-    # element i of t is the number of elements sharing rank i
-    _, t = np.unique(ranks, return_counts=True, axis=-1)
-    return (t**3 - t).sum(axis=-1)
-
-
-def _get_mwu_z(U, n1, n2, ranks, axis=0, continuity=True):
+def _get_mwu_z(U, n1, n2, t, axis=0, continuity=True):
     '''Standardized MWU statistic'''
     # Follows mannwhitneyu [2]
     mu = n1 * n2 / 2
     n = n1 + n2
 
-    # Tie correction according to [2]
-    tie_term = np.apply_along_axis(_tie_term, -1, ranks)
+    # Tie correction according to [2], "Normal approximation and tie correction"
+    # "A more computationally-efficient form..."
+    tie_term = (t**3 - t).sum(axis=-1)
     s = np.sqrt(n1*n2/12 * ((n + 1) - tie_term/(n*(n-1))))
 
-    # equivalent to using scipy.stats.tiecorrect
-    # T = np.apply_along_axis(stats.tiecorrect, -1, ranks)
-    # s = np.sqrt(T * n1 * n2 * (n1+n2+1) / 12.0)
-
     numerator = U - mu
 
     # Continuity correction.
@@ -223,13 +214,7 @@ def _mwu_input_validation(x, y, use_continuity, alternative, axis, method):
     return x, y, use_continuity, alternative, axis_int, method
 
 
-def _tie_check(xy):
-    """Find any ties in data"""
-    _, t = np.unique(xy, return_counts=True, axis=-1)
-    return np.any(t != 1)
-
-
-def _mwu_choose_method(n1, n2, xy, method):
+def _mwu_choose_method(n1, n2, ties, method):
     """Choose method 'asymptotic' or 'exact' depending on input size, ties"""
 
     # if both inputs are large, asymptotic is OK
@@ -237,7 +222,7 @@ def _mwu_choose_method(n1, n2, xy, method):
         return "asymptotic"
 
     # if there are any ties, asymptotic is preferred
-    if np.apply_along_axis(_tie_check, -1, xy).any():
+    if ties:
         return "asymptotic"
 
     return "exact"
@@ -471,14 +456,14 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
 
     n1, n2 = x.shape[-1], y.shape[-1]
 
-    if method == "auto":
-        method = _mwu_choose_method(n1, n2, xy, method)
-
     # Follows [2]
-    ranks = stats.rankdata(xy, axis=-1)  # method 2, step 1
-    R1 = ranks[..., :n1].sum(axis=-1)    # method 2, step 2
-    U1 = R1 - n1*(n1+1)/2                # method 2, step 3
-    U2 = n1 * n2 - U1                    # as U1 + U2 = n1 * n2
+    ranks, t = _rankdata(xy, 'average', return_ties=True)  # method 2, step 1
+    R1 = ranks[..., :n1].sum(axis=-1)                      # method 2, step 2
+    U1 = R1 - n1*(n1+1)/2                                  # method 2, step 3
+    U2 = n1 * n2 - U1                                      # as U1 + U2 = n1 * n2
+
+    if method == "auto":
+        method = _mwu_choose_method(n1, n2, np.any(t > 1), method)
 
     if alternative == "greater":
         U, f = U1, 1  # U is the statistic to use for p-value, f is a factor
@@ -490,7 +475,7 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
     if method == "exact":
         p = _mwu_state.sf(U.astype(int), n1, n2)
     elif method == "asymptotic":
-        z = _get_mwu_z(U, n1, n2, ranks, continuity=use_continuity)
+        z = _get_mwu_z(U, n1, n2, t, continuity=use_continuity)
         p = stats.norm.sf(z)
     p *= f
 


Code review comments:
diff_hunk: @@ -471,14 +456,14 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
 
     n1, n2 = x.shape[-1], y.shape[-1]
 
-    if method == "auto":
-        method = _mwu_choose_method(n1, n2, xy, method)
-
     # Follows [2]
-    ranks = stats.rankdata(xy, axis=-1)  # method 2, step 1
-    R1 = ranks[..., :n1].sum(axis=-1)    # method 2, step 2
-    U1 = R1 - n1*(n1+1)/2                # method 2, step 3
-    U2 = n1 * n2 - U1                    # as U1 + U2 = n1 * n2
+    ranks, t = _rankdata(xy, 'average', return_ties=True)  # method 2, step 1
+    R1 = ranks[..., :n1].sum(axis=-1)                      # method 2, step 2
+    U1 = R1 - n1*(n1+1)/2                                  # method 2, step 3
+    U2 = n1 * n2 - U1                                      # as U1 + U2 = n1 * n2
+
+    if method == "auto":
+        method = _mwu_choose_method(n1, n2, np.any(t > 1), method)
path: scipy/stats/_mannwhitneyu.py
text: @user1:
I think this two lines 465 and 466 can be moved to before line 475 for better readability because the `method` is used from line 475.

@author:
Ok. Will do.

diff_hunk: @@ -223,21 +214,15 @@ def _mwu_input_validation(x, y, use_continuity, alternative, axis, method):
     return x, y, use_continuity, alternative, axis_int, method
 
 
-def _tie_check(xy):
-    """Find any ties in data"""
-    _, t = np.unique(xy, return_counts=True, axis=-1)
-    return np.any(t != 1)
-
-
-def _mwu_choose_method(n1, n2, xy, method):
+def _mwu_choose_method(n1, n2, ties, method):
path: scipy/stats/_mannwhitneyu.py
text: @user1:
This function does not use the argument `method`. So, this can be removed?

@author:
I suppose so. I guess it was like that before. I'll do that.