scipy中Mann-Whitney U检验大规模计算的性能优化建议

背景 #

11年后，我饶有兴趣地阅读了一个Scipy相关的issue，编写了如下代码：

import numpy as np
from scipy.stats import mannwhitneyu

from tqdm import tqdm

# Basics
runs = 10000
scores = 100000
states = 1000

# Random scores
score = np.random.rand(runs * scores).reshape(runs, scores)
score.shape

# Random states
state = np.random.choice(
    [True, False],
    p=[0.01, 0.99],  # Roughly 1% True, 99% False (on average)
    size=scores * states,
).reshape(states, scores)
state.shape

muh_ps = np.zeros((runs, states))

for i in tqdm(range(runs), desc="Runs", total=runs):
    sc = score[i]

    for st in tqdm(range(states), desc="States", total=states):
        rp = sc[state[st]]
        rn = sc[~state[st]]

        u_statistic, u_pvalue = mannwhitneyu(rp, rn, alternative="two-sided")

        muh_ps[i, st] = u_pvalue

print(muh_ps)

当前使用的Scipy版本：

$ pip freeze | grep scipy
scipy==1.14.1

现有实现的性能问题 #

我认为不应采用双重循环的方式，但不确定正确的实现方法。
尝试了上述“朴素”实现，单轮速度尚可：

test_man.py 
States: 100%|███████████████████████████████████████| 1000/1000 [00:38<00:00, 25.86it/s]
States: 100%|███████████████████████████████████████| 1000/1000 [00:38<00:00, 25.74it/s]
States: 100%|███████████████████████████████████████| 1000/1000 [00:39<00:00, 25.36it/s]
States: 100%|███████████████████████████████████████| 1000/1000 [00:37<00:00, 26.43it/s]
States: 100%|███████████████████████████████████████| 1000/1000 [00:38<00:00, 26.14it/s]
Runs:   0%|                                         |   5/10000 [03:13<106:48:42, 38.47s/it]
States:  83%|██████████████████████████████████▌    |  827/1000 [00:27<00:06, 25.06it/s]

但按此速度完成全部计算需约5天。

优化请求 #

我尝试向mannwhitneyu传入布尔矩阵但未成功，考虑改用CuPy但需要指导，恳请提供优化方向及实现建议。

内容的提问来源于stack exchange，提问作者Dan Bolfter