使用 Cython 比较 scipy.stats.t.sf 与 GSL

Question

我想计算一个大型 2D numpy t 值数组的 p 值。但是，这需要很长时间，我想提高它的速度。我尝试使用 GSL。虽然单个 gsl_cdf_tdist_P 比 scipy.stats.t.sf 快很多，但在遍历 ndarray 时，过程非常缓慢。我想帮助改善这一点。请参阅下面的代码。

GSL_Test.pyx

import cython
cimport cython

import numpy
cimport numpy
from cython_gsl cimport *

DTYPE = numpy.float32
ctypedef numpy.float32_t DTYPE_t

cdef get_gsl_p(double t, double nu):
    return (1 - gsl_cdf_tdist_P(t, nu)) * 2

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef get_gsl_p_for_2D_matrix(numpy.ndarray[DTYPE_t, ndim=2] t_matrix, int n):

    cdef unsigned int rows = t_matrix.shape[0]
    cdef numpy.ndarray[DTYPE_t, ndim=2] out = numpy.zeros((rows, rows), dtype='float32')
    cdef unsigned int row, col

    for row in range(rows):
        for col in range(rows):
            out[row, col] = get_gsl_p(t_matrix[row, col], n-2)
    return out

def get_gsl_p_for_2D_matrix_def(numpy.ndarray[DTYPE_t, ndim=2] t_matrix, int n):
    return get_gsl_p_for_2D_matrix(t_matrix, n)

ipython

import GSL_Test
import numpy
import scipy.stats
a = numpy.random.rand(3544, 3544).astype('float32')
%timeit -n 1 GSL_Test.get_gsl_p_for_2D_matrix(a, 25)
1 loop, best of 3: 7.87 s per loop
%timeit -n 1 scipy.stats.t.sf(a, 25)*2
1 loop, best of 3: 4.66 s per loop

更新： 添加 cdef 声明我能够减少计算时间但仍不低于 scipy。我修改了代码以包含 cdef 声明。

%timeit -n 1 GSL_Test.get_gsl_p_for_2D_matrix_def(a, 25)
1 loop, best of 3: 6.73 s per loop

Answer 1

通过使用原始特殊函数而不是 stats.t.sf，您可以获得一些原始性能的小幅提升。查看源代码，您会发现 (https://github.com/scipy/scipy/blob/master/scipy/stats/_continuous_distns.py#L3849)

def _sf(self, x, df):
      return sc.stdtr(df, -x)

这样就可以直接使用stdtr:

np.random.seed(1234)
x = np.random.random((3740, 374))
t1 = stats.t.sf(x, 25)
t2 = stdtr(25, -x)

1 loop, best of 3: 653 ms per loop
1 loop, best of 3: 562 ms per loop

如果您确实使用了 cython，类型化的 memoryview 语法通常会比旧的 ndarray 语法提供更快的代码：

from scipy.special.cython_special cimport stdtr
from numpy cimport npy_intp
import numpy as np

def tsf(double [:, ::1] x, int df=25):
    cdef double[:, ::1] out = np.empty_like(x)
    cdef npy_intp i, j
    cdef double tmp, xx

    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            xx = x[i, j]
            out[i, j] = stdtr(df, -xx)

    return np.asarray(out)

这里我也使用了cython_special接口，只有scipy(http://scipy.github.io/devdocs/special.cython_special.html#module-scipy.special.cython_special)的dev版本才有，但是如果你愿意，你可以使用GSL。

最后，如果您怀疑迭代中存在瓶颈，请不要忘记检查 cython -a 的输出以查看热循环中是否存在一些 python 开销。

使用 Cython 比较 scipy.stats.t.sf 与 GSL

Comparing scipy.stats.t.sf vs GSL using Cython

python

numpy

cython

scipy

gsl