如何加速scipy.stats.hypergeom超几何生存函数的大规模计算

可用的加速方案 #

方案1：向量化调用替代逐次循环 #

Scipy的hypergeom.sf原生支持数组作为输入，你现在逐次传入单个值的写法浪费了大量的函数调用开销和内部初始化成本，直接批量传入所有参数即可，改动成本极低：

import numpy as np
from scipy import stats

fftot = 40071
ffb = 375

# 先把输入转成numpy数组
object_ontology_counts_arr = np.array(object_ontology_counts)
oe_overlaps_arr = np.array(oe_overlaps)

# 批量计算所有结果
results = np.where(oe_overlaps_arr <= 0, 1.0,
                   stats.hypergeom.sf(oe_overlaps_arr - 1, fftot, object_ontology_counts_arr, ffb))
# 若需要保持和原代码一致的列表格式可加以下转换
results = results.reshape(-1, 1).tolist()

该方案无需额外依赖，通常就能获得10~50倍的提速。

方案2：用Numba JIT编译计算逻辑 #

你场景中的fftot和ffb两个参数是固定常量，用Numba对计算逻辑做JIT编译后，性能可以接近原生C水平：
首先安装依赖：pip install numba
代码示例：

import numpy as np
from numba import jit

fftot = 40071
ffb = 375

@jit(nopython=True)
def calculate_single(object_ontology_count, overlap):
    if overlap <= 0:
        return 1.0
    return stats.hypergeom.sf(overlap - 1, fftot, object_ontology_count, ffb)

# 生成向量化计算函数
calculate_vec = np.vectorize(calculate_single)
results = calculate_vec(np.array(object_ontology_counts), np.array(oe_overlaps))

该方案可以获得50~100倍的提速，操作难度很低。

方案3：预计算缓存结果 #

你的两个输入参数取值范围非常有限：overlap最大不超过375，object_ontology_count最大不超过40071，所有可能的参数组合最多只有40071 * 375 ≈ 1500万种，远小于你需要计算的9亿次。你可以提前把所有可能的参数对应的结果计算好存为二维数组，后续直接查表即可，查询速度可以达到每秒数亿次：

import numpy as np
from scipy import stats

fftot = 40071
ffb = 375

max_ob_count = fftot
max_overlap = ffb

# 预构建缓存表
cache = np.zeros((max_ob_count + 1, max_overlap + 1), dtype=np.float64)
for ob_count in range(1, max_ob_count + 1):
    for overlap in range(0, max_overlap + 1):
        if overlap <= 0:
            cache[ob_count][overlap] = 1.0
        else:
            cache[ob_count][overlap] = stats.hypergeom.sf(overlap - 1, fftot, ob_count, ffb)

# 后续9亿次计算直接查表即可
def get_hypergeom_val(ob_count, overlap):
    return cache[ob_count][min(overlap, max_overlap)]

该方案是性能最高的，预计算仅需要几分钟，后续查询几乎没有开销。

方案4：调用R的原生实现 #

如果接受安装R环境，可以用rpy2调用R原生的phyper函数，性能和R原生一致，比Scipy实现快接近100倍，仅需简单封装即可使用。

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