Repository: scipy/scipy
Pull request: ENH: stats.mannwhitneyu: speed improvement, memory reduction, and `PermutationMethod` support
#### Reference issue
Closes gh-19692

#### What does this implement/fix?
gh-19692 reported high memory consumption in the `stats.mannwhitneyu` p-value calculation. Several improvements were suggested in https://github.com/scipy/scipy/issues/19692#issuecomment-1863311504. This implements most of those improvement ideas.

#### Additional information
Before this PR, both of these calls would have been slow, and memory consumption would have increased during the second call. Now, second call is very fast, and it doesn't increase memory usage.

```python3
import numpy as np
from scipy import stats
rng = np.random.default_rng(2549824598234528)
x = rng.random(size=30)
y = rng.random(size=120)
stats.mannwhitneyu(x, y, method='exact')
# second call is fast because cache can be reused even when
# sample sizes are reversed
stats.mannwhitneyu(y, x, method='exact')
```

Before this PR, this call would have been very slow because the p-value computation required summing all elements of the null distribution PDF. Now, by exploiting the symmetry of the null distribution, only one element of the null distribution PDF needs to be calculated.

```python3
import numpy as np
from scipy import stats
rng = np.random.default_rng(2549824598234528)
x = rng.random(size=30)
y = rng.random(size=120)
stats.mannwhitneyu(0*x, y, method='exact', alternative='greater')
```

Finally, this PR adds support for `method` being an instance of `PermutationMethod`. This can p-values to be estimated accurately in cases where both `asymptotic` and `exact` would not be acceptable (e.g. small samples with ties, one sample very large with the other very small):
```python3
import numpy as np
from scipy import stats
rng = np.random.default_rng(2549824598234528)
rng = np.random.default_rng(2549824598234528)
x = rng.random(size=5)
y = rng.random(size=6)
stats.mannwhitneyu(x, y, method='exact')
# MannwhitneyuResult(statistic=9.0, pvalue=0.32900432900432897)
stats.mannwhitneyu(x, y, method=stats.PermutationMethod())
# MannwhitneyuResult(statistic=9.0, pvalue=0.329004329004329)

x[1] = x[2]  # introduce a tie
stats.mannwhitneyu(x, y, method='exact')
# MannwhitneyuResult(statistic=12.0, pvalue=0.6623376623376623)
stats.mannwhitneyu(x, y, method=stats.PermutationMethod())
# MannwhitneyuResult(statistic=12.0, pvalue=0.6277056277056277)
```

Problem statement:
BUG: memory issues in mannwhitneyu with unbalanced data
### Describe your issue.

One of our pipelines calls the `mannwhitneyu` method using the "two-sided" alternative.  We've noticed that the method can consume a substantial amount of RAM. For example, in a call with a total of 17,002 values the `mannwhitneyu` call peaks at nearly 24GB of RAM.

There appears to be a memory leak as well.  Sometimes a single call such as the one above will finish despite the huge memory footprint.  However if our pipeline calls it several times, eventually the process fails with an out-of-memory error.

Full disclosure, some of our users provide data that becomes wildly unbalanced, e.g., 2 values in one array and 17,000 in the other (see example code).  I don't know what impact that will have on the method.  We're investigating ways to prevent such unbalanced calls on our end.  However, the issue remains that even with more balanced data this method uses far more memory than one would expect.


### Reproducing Code Example

```python
"""
Example of scipy.stats.mannwhitneyu memory issues
"""
import numpy
from scipy.stats import mannwhitneyu
A = [-0.0123, -0.1023]
B = numpy.random.uniform(low=-1, high=1, size=17000)
# On an M1 Mac with 24GB allocated to Docker containers, this will take
# about 5 minutes before it runs out of memory inside a container:
result = mannwhitneyu(A, B, alternative="two-sided")
```


### Error message

```shell
Killed
```


### SciPy/NumPy/Python version and system information

```shell
Rocky Linux 9
confirmed on scipy version 1.10.0 and 1.11.4
```

Hints from issue discussion:
Like many SciPy functions, `mannwhitneyu` doesn't consider the memory limitations of your machine when deciding which method to use. If either of the two samples has fewer than 8 observations, and `method='auto'` (the default), it still chooses the exact method, which can consume large amounts of memory. You can select `method='asymptotic'` to dramatically reduce memory and time requirements, but the p-values won't necessarily be accurate. 
```python3
result = mannwhitneyu(A, B, alternative="two-sided", method='asymptotic')
# MannwhitneyuResult(statistic=16159.0, pvalue=0.9036158130132622)
```
Your best bet is probably to estimate the p-value using a randomized permutation test. Some of the last digits will fluctuate between individual runs, but in the long term it almost exactly controls the false positive error rate. If that sounds acceptable, here is some code.

<details>
Let me double check the alternative behavior before anyone relies on this.

```python3
from scipy.stats import mannwhitneyu
import numpy as np
rng = np.random.default_rng(345823468354843567)
A = [-0.0123, -0.1023]
B = rng.uniform(low=-1, high=1, size=17000)
result = mannwhitneyu(A, B, alternative="two-sided", method='asymptotic')

from scipy import stats

def statistic(x, y):
  return stats.mannwhitneyu(x, y, alternative="two-sided", 
                            method='asymptotic').statistic

res = stats.permutation_test((A, B), statistic, random_state=rng, 
                             batch=100, n_resamples=9999)
print(result.pvalue, res.pvalue)
# 0.8792445100889502 0.8856
```

</details>

> There appears to be a memory leak as well. 

There is an object that persists in memory (before and) after such a call to make future calls with samples of the same size faster. Calling with larger samples will increase the size of the single object. So, suppose you were to make a call with 2 and 17000 observations in `x` and `y`, respectively, then you were to call it with 100 and 100 observations; IIRC this might expand the tables to a size you'd get if you had called it with 100 and 17000 observations. So it's an unintended consequence of the strategy for improving performance with moderately-sized samples. Could you provide an example of this sort?
Many thanks for your response, and for helping me understand the increasing memory issue.  In our case, we may get every one of the combinations of 1, 2, ..., 8 vs. 17,000, in addition to many others, so I think that explains the big memory hit after ~400 calls or so.  I will try to build an example script in the next day or so to share with you.

Thanks also for the permutation test example.  We did find the `method="asymptotic"` fix yesterday, but it's entirely possible we'll see misleading p-values with those small `A` samples.
I've finally had a moment to create an example (attached below) where the memory footprint grows gradually until it consumes all available memory, at least on a system (Docker image) limited to 16GB.  I believe the same could happen for systems with up to 32GB.

I added some rudimentary timing and ran each `mannwhitneyu` call twice to see the impact of the optimization.  Surprisingly, there appears to be no difference between the two calls.  Perhaps this isn't a fair test of the optimization, but I thought it worth pointing out.
hi @mrogers-phd, please upload the data to a source where the text source can be read, `zip` files are deemed a security risk
Argh.  When I tried to upload just the python script, github asked for a different file type such as `.zip`.  Probably just as simple to show the source itself:
```
#! /usr/bin/env python3
import numpy
import time
from scipy.stats import mannwhitneyu

# Docker image with 16GB memory (host computer has 32GB)
# A (size 2) vs B (size 5000)
#   p-value:  0.77147 (compute time 31.5 secs)
#   2nd call: 0.77147 (compute time 30.8 secs)
# A (size 3) vs B (size 5000)
#   p-value:  0.99235 (compute time 72.7 secs)
#   2nd call: 0.99235 (compute time 73.1 secs)
# A (size 4) vs B (size 5000)
#   p-value:  0.4283 (compute time 92.2 secs)
#   2nd call: 0.4283 (compute time 91.6 secs)
# A (size 5) vs B (size 5000)
# Killed

# Large baseline sample for all comparisons
rng = numpy.random.default_rng(12345)
B = rng.uniform(low=-1, high=1, size=5000)

# Source list for small samples:
A_source = rng.uniform(low=-1, high=1, size=10)

for i in range(1, len(A_source)):
    # Create subsets for sizes 2, 3, ...
    A = A_source[:i+1]
    print(f"A (size {len(A)}) vs B (size {len(B)})")

    start = time.time()
    result = mannwhitneyu(A, B, alternative="two-sided")
    elapsed = time.time() - start
    print(f"  p-value:  {result.pvalue:.5g} (compute time {elapsed:.1f} secs)")

    start = time.time()
    repeated = mannwhitneyu(A, B, alternative="two-sided")
    elapsed = time.time() - start
    print(f"  2nd call: {repeated.pvalue:.5g} (compute time {elapsed:.1f} secs)")
```
Hmm, something is not working as expected. I'll take a look.
Minor copy/paste error there on the last line  (was `result.pvalue`, should have been `repeated.pvalue`) that I've fixed above, but shouldn't have had an impact on the overall results.
I will look more carefully, but I see that what I've written before doesn't apply - when the size of either sample is greater than 500, we switch from a recursive implementation to an iterative implementation, which doesn't use the table I was referring to. Apparently this means everything is calculated from scratch, so you don't get the speedup. However, this also means that a growing table doesn't explain why you are out of memory on a second call. 

Instead, I notice that the number of calculations that have to be performed and stored depends on the value of the statistic, so that is my new guess for why you might succeed in one call but not the next. It looks like worst case memory usage is $O(m^2n^2)$ in the recursive case, and if the iterative version is similar, the story might be better or worse because it uses a dictionary instead of an array for that storage.

There are a few potential improvements that catch my eye.
- Even during initial development, it was noted that the symmetry of the null distribution could be exploited more fully. This could potentially improve the memory usage and speed if the value of the statistic is in the tails of the distribution. 
- I don't think this helps with the iterative method since it doesn't reuse tabulated values, but if the samples are very unbalanced and the larger one can sometimes be first and sometimes be second, there is some symmetry between the two samples that can be exploited. For example, for the two-sided null hypothesis, it shouldn't matter which sample is passed first, but currently they are computed differently. If the p-values were computed with the larger sample always first (or second), we could avoid repeating work and unnecessarily increasing table size.
- the iterative method would see a speed improvement if we were to reuse tabulated values between runs.

But these probably won't make or break things in your use case unless you know your statistic value is never close to the mean. What sort of p-values do you tend to see? Do you occasionally see both moderate and large p-values?

I think the best bet is to add support for a randomized permutation test by accepting a [`PermutationMethod`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.PermutationMethod.html#scipy.stats.PermutationMethod) as the `method` parameter.
In our process we can get p-values anywhere in the range (0,1) and our significance threshold is typically p < 0.01.  From what we've observed so far, the asymptotic method works well, though in rare cases it can produce a significant p-value that is noticeably higher than the exact method (e.g., 1e-05 instead of 1e-08).  That may become an issue when we later apply multi-test correction.
Huh, well the p-values provided by the randomized permutation test are not ever going to be less than the reciprocal of the number of resamples, so you'd have the same problem in that case.

Whether the implementation uses iteration or explicit recursion, the algorithm is based on a recursive formula. Currently, every evaluation of the function that is being recursed is cached. Perhaps some analysis of the algorithm would reveal ways of reducing this memory consumption. (Of course, maybe there are other, inherently more memory-efficient algorithms.)

I'm not able to look deeply into that right now, but perhaps there is a different solution in the short term: If you don't need super accurate p-values above your significance threshold but do need the precision of the exact method below the threshold, then you could use the approximate p-value as an estimate and refine it using the exact method only if the p-value is small. Theoretically, at least, the memory (and time) required should be less for small p-values.

Commit message: MAINT: stats.mannwhitneyu: minor adjustments

Patch to review:
diff --git a/scipy/stats/_mannwhitneyu.py b/scipy/stats/_mannwhitneyu.py
index 342eccbacf87..99fb78cad3e4 100644
--- a/scipy/stats/_mannwhitneyu.py
+++ b/scipy/stats/_mannwhitneyu.py
@@ -26,6 +26,14 @@ def __init__(self):
         self._recursive = None
 
     def pmf(self, k, m, n):
+
+        # In practice, `pmf` is never called with k > m*n/2.
+        # If it were, we'd exploit symmetry here:
+        # k = np.array(k, copy=True)
+        # k2 = m*n - k
+        # i = k2 < k
+        # k[i] = k2[i]
+
         if (self._recursive is None and m <= 500 and n <= 500
                 or self._recursive):
             return self.pmf_recursive(k, m, n)
@@ -51,22 +59,30 @@ def pmf_iterative(self, k, m, n):
 
     def cdf(self, k, m, n):
         '''Cumulative distribution function'''
-        # We could use the fact that the distribution is symmetric to avoid
-        # summing more than m*n/2 terms, but it might not be worth the
-        # overhead. Let's leave that to an improvement.
+
+        # In practice, `cdf` is never called with k > m*n/2.
+        # If it were, we'd exploit symmetry here rather than in `sf`
         pmfs = self.pmf(np.arange(0, np.max(k) + 1), m, n)
         cdfs = np.cumsum(pmfs)
         return cdfs[k]
 
     def sf(self, k, m, n):
         '''Survival function'''
-        # Use the fact that the distribution is symmetric; i.e.
-        # _f(m, n, m*n-k) = _f(m, n, k), and sum from the left
-        k = m*n - k
         # Note that both CDF and SF include the PMF at k. The p-value is
         # calculated from the SF and should include the mass at k, so this
         # is desirable
-        return self.cdf(k, m, n)
+
+        # Use the fact that the distribution is symmetric; i.e.
+        # _f(m, n, m*n-k) = _f(m, n, k), and sum from the left
+        kc = np.asarray(m*n - k)  # complement of k
+        i = k < kc
+        if np.any(i):
+            kc[i] = k[i]
+            cdfs = np.asarray(self.cdf(kc, m, n))
+            cdfs[i] = 1. - cdfs[i] + self.pmf(kc[i], m, n)
+        else:
+            cdfs = np.asarray(self.cdf(kc, m, n))
+        return cdfs[()]
 
     def _resize_fmnks(self, m, n, k):
         '''If necessary, expand the array that remembers PMF values'''
@@ -215,10 +231,11 @@ def _mwu_input_validation(x, y, use_continuity, alternative, axis, method):
     if axis != axis_int:
         raise ValueError('`axis` must be an integer.')
 
-    methods = {"asymptotic", "exact", "auto"}
-    method = method.lower()
-    if method not in methods:
-        raise ValueError(f'`method` must be one of {methods}.')
+    if not isinstance(method, stats.PermutationMethod):
+        methods = {"asymptotic", "exact", "auto"}
+        method = method.lower()
+        if method not in methods:
+            raise ValueError(f'`method` must be one of {methods}.')
 
     return x, y, use_continuity, alternative, axis_int, method
 
@@ -291,7 +308,7 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
         see [5] section 5.1.
     axis : int, optional
         Axis along which to perform the test. Default is 0.
-    method : {'auto', 'asymptotic', 'exact'}, optional
+    method : {'auto', 'asymptotic', 'exact'} or `stats.PermutationMethod`, optional
         Selects the method used to calculate the *p*-value.
         Default is 'auto'. The following options are available.
 
@@ -304,6 +321,10 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
           is less than or equal to 8 and there are no ties;
           chooses ``'asymptotic'`` otherwise.
 
+        `method` may also be an instance of `stats.PermutationMethod`, in which
+        case the p-value is computed using `scipy.stats.permutation_test` with
+        the provided configuration options and other appropriate settings.
+
     Returns
     -------
     res : MannwhitneyuResult
@@ -329,7 +350,9 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
     relation originally proposed in [1]_ as it is described in [3]_.
     Note that the exact method is *not* corrected for ties, but
     `mannwhitneyu` will not raise errors or warnings if there are ties in the
-    data.
+    data. If there are ties and either samples is small (fewer than ~10
+    observations), consider passing an instance of `stats.PermutationMethod`
+    as the `method` to perform a permutation test.
 
     The Mann-Whitney U test is a non-parametric version of the t-test for
     independent samples. When the means of samples from the populations
@@ -488,10 +511,21 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
         U, f = np.maximum(U1, U2), 2  # multiply SF by two for two-sided test
 
     if method == "exact":
-        p = _mwu_state.sf(U.astype(int), n1, n2)
+        p = _mwu_state.sf(U.astype(int), min(n1, n2), max(n1, n2))
     elif method == "asymptotic":
         z = _get_mwu_z(U, n1, n2, ranks, continuity=use_continuity)
         p = stats.norm.sf(z)
+    else:
+        def statistic(x, y, axis):
+            return mannwhitneyu(x, y, use_continuity=use_continuity,
+                                alternative=alternative, axis=axis,
+                                method="asymptotic").statistic
+
+        res = stats.permutation_test((x, y), statistic, axis=axis,
+                                     **method._asdict(), alternative=alternative)
+        p = res.pvalue
+        f = 1
+
     p *= f
 
     # Ensure that test statistic is not greater than 1
diff --git a/scipy/stats/tests/test_hypotests.py b/scipy/stats/tests/test_hypotests.py
index 52b316e07c1b..c6f9adbcdc99 100644
--- a/scipy/stats/tests/test_hypotests.py
+++ b/scipy/stats/tests/test_hypotests.py
@@ -609,6 +609,43 @@ def test_gh_2118(self, x, y, alternative, expected):
                            method="asymptotic")
         assert_allclose(res, expected, rtol=1e-12)
 
+    def test_gh19692_smaller_table(self):
+        # In gh-19692, we noted that the shape of the cache used in calculating
+        # p-values was dependent on the order of the inputs because the sample
+        # sizes n1 and n2 changed. This was indicative of unnecessary cache
+        # growth and redundant calculation. Check that this is resolved.
+        rng = np.random.default_rng(7600451795963068007)
+        x = rng.random(size=5)
+        y = rng.random(size=11)
+        _mwu_state._fmnks = -np.ones((1, 1, 1))  # reset cache
+        stats.mannwhitneyu(x, y, method='exact')
+        shape = _mwu_state._fmnks.shape
+        assert shape[0] <= 6 and shape[1] <= 12  # one more than sizes
+        stats.mannwhitneyu(y, x, method='exact')
+        assert shape == _mwu_state._fmnks.shape  # unchanged when sizes are reversed
+
+        # Also, we weren't exploiting the symmmetry of the null distribution
+        # to its full potential. Ensure that the null distribution is not
+        # evaluated explicitly for `k > m*n/2`.
+        _mwu_state._fmnks = -np.ones((1, 1, 1))  # reset cache
+        stats.mannwhitneyu(x, 0*y, method='exact', alternative='greater')
+        shape = _mwu_state._fmnks.shape
+        assert shape[-1] == 1  # k is smallest possible
+        stats.mannwhitneyu(0*x, y, method='exact', alternative='greater')
+        assert shape == _mwu_state._fmnks.shape
+
+    @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided'])
+    def test_permutation_method(self, alternative):
+        rng = np.random.default_rng(7600451795963068007)
+        x = rng.random(size=(2, 5))
+        y = rng.random(size=(2, 6))
+        res = stats.mannwhitneyu(x, y, method=stats.PermutationMethod(),
+                                 alternative=alternative, axis=1)
+        res2 = stats.mannwhitneyu(x, y, method='exact',
+                                  alternative=alternative, axis=1)
+        assert_allclose(res.statistic, res2.statistic, rtol=1e-15)
+        assert_allclose(res.pvalue, res2.pvalue, rtol=1e-15)
+
     def teardown_method(self):
         _mwu_state._recursive = None
 


Code review comments:
diff_hunk: @@ -304,6 +321,10 @@ def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided",
           is less than or equal to 8 and there are no ties;
           chooses ``'asymptotic'`` otherwise.
 
+        `method` may also be an instance of `stats.PermutationMethod`, in which
path: scipy/stats/_mannwhitneyu.py
text: @user1:
```suggestion
        * `scipy.stats.PermutationMethod`: `method` may also be an instance of `scipy.stats.PermutationMethod`, in which
```
I think Permutation method instance should be one of bullet item and making a doc link of `scipy.stats.PermutationMethod` is useful for readers. 
(currently this is not a link
![IMG_1491](https://github.com/scipy/scipy/assets/3813847/1f37f428-bc97-4503-bb75-a62dceed4ac2))