Repository: scipy/scipy
Pull request: ENH: stats.mannwhitneyu: add iterative implementation
#### Reference issue
Closes gh-14622
Supersedes gh-14854

#### What does this implement/fix?
gh-14622 reported that `mannwhitenyu`'s implementation fails for sufficiently large samples because it reached the maximum recursion depth. This PR adds an iterative implementation that kicks in when either of the sample sizes exceeds 500. 

#### Additional information
I wrote this because converting the recursive implementation to an iterative implementation was easier for me than finding and correcting the bug that stalled gh-14854. This is also slightly faster than the version in gh-14854 for the cases tested, and it lends itself well to being pythranized. 


Problem statement:
BUG: (sort of) mannwhitneyu hits max recursion limit with imbalanced input sizes due to selection of 'exact' method; could method selection depend on max(n1,n2)?
### Reproducing Code Example

```python
from scipy.stats import mannwhitneyu
import numpy as np
a = np.random.randn(5)
b = np.random.randn(5000)
mannwhitneyu(a,b)
```


### Error message

```shell
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 414, in mannwhitneyu
    p = _mwu_state.sf(U.astype(int), n1, n2)
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 53, in sf
    return self.cdf(k, m, n)
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 41, in cdf
    pmfs = self.pmf(np.arange(0, np.max(k) + 1), m, n)
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 33, in pmf
    self._f(m, n, i)
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 81, in _f
    fmnk = self._f(m-1, n, k-n) + self._f(m, n-1, k)
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 81, in _f
    fmnk = self._f(m-1, n, k-n) + self._f(m, n-1, k)
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 81, in _f
    fmnk = self._f(m-1, n, k-n) + self._f(m, n-1, k)
  [Previous line repeated 991 more times]
  File "/Users/pbradley/opt/anaconda3/envs/scanpy_piptest5/lib/python3.8/site-packages/scipy/stats/_mannwhitneyu.py", line 71, in _f
    if k < 0 or m < 0 or n < 0 or k > m*n:
RecursionError: maximum recursion depth exceeded in comparison
```


### SciPy/NumPy/Python version information

1.7.1 1.21.2 sys.version_info(major=3, minor=8, micro=11, releaselevel='final', serial=0)

Hints from issue discussion:
Hi there, sorry if this isn't a BUG per se, but it did break our code when we updated to the newest scipy. Of course, there's an easy work-around, but I wonder if the method selection logic could take into account the sizes of both inputs.
Good idea. 
You can add `method='asymptotic'` to get the same p-values you were getting before.
Thanks for reporting!
Thanks, yes, that's what I am doing now (specifying method). It's just a tiny bit tricky since older mannwhitneyu code (which we still need to account for) doesn't know about method as an argument, so I have to do a little check first.

To a contributor interested in this problem: as suggested above, the quick fix would be to figure out some additional conditions for choosing the asymptotic method and add them [here](https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/stats/_mannwhitneyu.py#L164).

Another possibility would be for [`_MWU`](https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/stats/_mannwhitneyu.py#L19) to remember its own stack.
@mdhaber I have an iterative solution that post-order traverses the stack (binary tree) with memoization. It is not affected by recursion limits. It uses a list as a stack. Solution is equivalent to your implementation but takes about twice the time as native recursion. I ran a suite of tests and the p-value is not entirely insensitive to switching to 'asymptotic' when at least one of n and m is small. Therefore, IMHO, the best course of action is to use a check on n and m and if one is <8 and the other is large (would hit the recursion limit), switch to the iterative solution (still 'exact') and issue a warning about possibly long run time. I don't think the iterative solution can be as fast as recursion. Would you be OK with iteration being a little slower than recursion?
@szkafander Thanks, but I (personally) would not be interested in making that change at this time. Of course, if others are interested, go for it; I have no particular say here - I just wouldn't be able to review it right now. In the meantime, I think gh-14835 is a good solution, considering the number of [other shortcomings in scipy.stats](https://github.com/scipy/scipy/labels/scipy.stats) that deserve attention. If you are interested in pursuing an iterative solution, though, I'd suggest trying to work it into the `_MWU` class, keeping in mind that some arguments can be arrays. 
@mdhaber I think I implemented the iterative solution that you outlined [here](https://github.com/scipy/scipy/issues/14622#issuecomment-902841703). Yes, my changes are in _MWU. I added an iterative version of _f. I am aware that 'k' is an iterable in some of the methods if that's what you mean.
OK, well, feel free to submit a PR. I might not be able to fully review for a while, but I'm curious what it looks like. 
@mdhaber sent [PR](https://github.com/scipy/scipy/pull/14854).
> In the meantime, I think [gh-14835](https://github.com/scipy/scipy/pull/14835) is a good solution

Oops, didn't see this before reviewing gh-14854. I think we can consider merging both gh-14854 and gh-14835 i.e. if the user really wants to use the `exact` method by explicitly passing the `exact` keyword, they can do so without hitting a recursion limit but if `method='auto'`, we can switch to the asymptotic method internally as we know it is fast and accurate for larger inputs(or throw an error as you suggested in gh-14835). Let me know what you think.

Commit message: ENH: stats.mannwhitenyu: enhance iterative exact method

Patch to review:
diff --git a/scipy/stats/_mannwhitneyu.py b/scipy/stats/_mannwhitneyu.py
index d1e2ef9f127c..eac9efb76a7b 100644
--- a/scipy/stats/_mannwhitneyu.py
+++ b/scipy/stats/_mannwhitneyu.py
@@ -23,15 +23,32 @@ class _MWU:
     def __init__(self):
         '''Minimal initializer'''
         self._fmnks = -np.ones((1, 1, 1))
+        self._recursive = None
 
     def pmf(self, k, m, n):
-        '''Probability mass function'''
+        if (self._recursive is None and m <= 500 and n <= 500
+                or self._recursive):
+            return self.pmf_recursive(k, m, n)
+        else:
+            return self.pmf_iterative(k, m, n)
+
+    def pmf_recursive(self, k, m, n):
+        '''Probability mass function, recursive version'''
+        self._fmnks = -np.ones((1, 1, 1))
         self._resize_fmnks(m, n, np.max(k))
         # could loop over just the unique elements, but probably not worth
         # the time to find them
         for i in np.ravel(k):
             self._f(m, n, i)
-        return self._fmnks[m, n, k] / special.binom(m + n, m)
+        return self._fmnks[(m, n, k)] / special.binom(m + n, m)
+
+    def pmf_iterative(self, k, m, n):
+        '''Probability mass function, iterative version'''
+        fmnks = {}
+        for i in np.ravel(k):
+            fmnks = _mwu_f_iterative(m, n, i, fmnks)
+        return (np.array([fmnks[(m, n, ki)] for ki in k])
+                / special.binom(m + n, m))
 
     def cdf(self, k, m, n):
         '''Cumulative distribution function'''
@@ -89,6 +106,58 @@ def _f(self, m, n, k):
 _mwu_state = _MWU()
 
 
+def _mwu_f_iterative(m, n, k, fmnks):
+    '''Iterative implementation of function of [3] Theorem 2.5'''
+
+    def _base_case(m, n, k):
+        '''Base cases from recursive versions'''
+
+        # if already calculated, return the value
+        if fmnks.get((m, n, k), -1) >= 0:
+            return fmnks[(m, n, k)]
+
+        # [3] Theorem 2.5 Line 1
+        elif k < 0 or m < 0 or n < 0 or k > m*n:
+            return 0
+
+        # [3] Theorem 2.5 Line 2
+        elif k == 0 and m >= 0 and n >= 0:
+            return 1
+
+        return None
+
+    stack = [(m, n, k)]
+    fmnk = None
+
+    while stack:
+        # Popping only if necessary would save a tiny bit of time, but NWI.
+        m, n, k = stack.pop()
+
+        # If we're at a base case, continue (stack unwinds)
+        fmnk = _base_case(m, n, k)
+        if fmnk is not None:
+            fmnks[(m, n, k)] = fmnk
+            continue
+
+        # If both terms are base cases, continue (stack unwinds)
+        f1 = _base_case(m-1, n, k-n)
+        f2 = _base_case(m, n-1, k)
+        if f1 is not None and f2 is not None:
+            # [3] Theorem 2.5 Line 3 / Equation 3
+            fmnk = f1 + f2
+            fmnks[(m, n, k)] = fmnk
+            continue
+
+        # recurse deeper
+        stack.append((m, n, k))
+        if f1 is None:
+            stack.append((m-1, n, k-n))
+        if f2 is None:
+            stack.append((m, n-1, k))
+
+    return fmnks
+
+
 def _tie_term(ranks):
     """Tie correction term"""
     # element i of t is the number of elements sharing rank i
diff --git a/scipy/stats/tests/test_hypotests.py b/scipy/stats/tests/test_hypotests.py
index 64e9bf3c8b95..8a6fda1abd15 100644
--- a/scipy/stats/tests/test_hypotests.py
+++ b/scipy/stats/tests/test_hypotests.py
@@ -172,6 +172,9 @@ def test_callable_cdf(self):
 
 
 class TestMannWhitneyU:
+    def setup(self):
+        _mwu_state._recursive = True
+
     # All magic numbers are from R wilcox.test unless otherwise specied
     # https://rdrr.io/r/stats/wilcox.test.html
 
@@ -616,6 +619,43 @@ def test_gh_2118(self, x, y, alternative, expected):
                            method="asymptotic")
         assert_allclose(res, expected, rtol=1e-12)
 
+    def teardown(self):
+        _mwu_state._recursive = None
+
+
+class TestMannWhitneyU_iterative(TestMannWhitneyU):
+    def setup(self):
+        _mwu_state._recursive = False
+
+    def teardown(self):
+        _mwu_state._recursive = None
+
+
+@pytest.mark.xslow
+def test_mann_whitney_u_switch():
+    # Check that mannwhiteneyu switches between recursive and iterative
+    # implementations at n = 500
+
+    # ensure that recursion is not enforced
+    _mwu_state._recursive = None
+    _mwu_state._fmnks = -np.ones((1, 1, 1))
+
+    rng = np.random.default_rng(9546146887652)
+    x = rng.random(5)
+
+    # use iterative algorithm because n > 500
+    y = rng.random(501)
+    stats.mannwhitneyu(x, y, method='exact')
+    # iterative algorithm doesn't modify _mwu_state._fmnks
+    assert np.all(_mwu_state._fmnks == -1)
+
+    # use recursive algorithm because n <= 500
+    y = rng.random(500)
+    stats.mannwhitneyu(x, y, method='exact')
+
+    # iterative algorithm has modified _mwu_state._fmnks
+    assert not np.all(_mwu_state._fmnks == -1)
+
 
 class TestSomersD(_TestPythranFunc):
     def setup_method(self):


Code review comments:
diff_hunk: @@ -23,15 +23,32 @@ class _MWU:
     def __init__(self):
         '''Minimal initializer'''
         self._fmnks = -np.ones((1, 1, 1))
+        self._recursive = None
 
     def pmf(self, k, m, n):
-        '''Probability mass function'''
+        if (self._recursive is None and m <= 500 and n <= 500
+                or self._recursive):
+            return self.pmf_recursive(k, m, n)
+        else:
+            return self.pmf_iterative(k, m, n)
+
+    def pmf_recursive(self, k, m, n):
+        '''Probability mass function, recursive version'''
+        self._fmnks = -np.ones((1, 1, 1))
         self._resize_fmnks(m, n, np.max(k))
         # could loop over just the unique elements, but probably not worth
         # the time to find them
         for i in np.ravel(k):
             self._f(m, n, i)
-        return self._fmnks[m, n, k] / special.binom(m + n, m)
+        return self._fmnks[(m, n, k)] / special.binom(m + n, m)
path: scipy/stats/_mannwhitneyu.py
text: @author:
```suggestion
        return self._fmnks[m, n, k] / special.binom(m + n, m)
```
This must have crept in when I was exploring use of `dict` to avoid running out of memory.

diff_hunk: @@ -23,15 +23,32 @@ class _MWU:
     def __init__(self):
         '''Minimal initializer'''
         self._fmnks = -np.ones((1, 1, 1))
+        self._recursive = None
 
     def pmf(self, k, m, n):
-        '''Probability mass function'''
+        if (self._recursive is None and m <= 500 and n <= 500
+                or self._recursive):
+            return self.pmf_recursive(k, m, n)
+        else:
+            return self.pmf_iterative(k, m, n)
+
+    def pmf_recursive(self, k, m, n):
+        '''Probability mass function, recursive version'''
+        self._fmnks = -np.ones((1, 1, 1))
path: scipy/stats/_mannwhitneyu.py
text: @author:
```suggestion
```
This was for testing.