高效计算列相关系数

Question

原始问题

我正在将大小为n的行P与大小为n×m的矩阵O的每一列进行相关。我编写了以下代码：

import numpy as np

def ColumnWiseCorrcoef(O, P):
    n = P.size
    DO = O - (np.sum(O, 0) / np.double(n))
    DP = P - (np.sum(P) / np.double(n))
    return np.dot(DP, DO) / np.sqrt(np.sum(DO ** 2, 0) * np.sum(DP ** 2))

它比朴素方法更有效：

def ColumnWiseCorrcoefNaive(O, P):
    return np.corrcoef(P,O.T)[0,1:O[0].size+1]

这是我在Intel Core上使用numpy-1.7.1-MKL得到的时间:

O = np.reshape(np.random.rand(100000), (1000,100))
P = np.random.rand(1000)

%timeit -n 1000 A = ColumnWiseCorrcoef(O, P)
1000 loops, best of 3: 787 us per loop
%timeit -n 1000 B = ColumnWiseCorrcoefNaive(O, P)
1000 loops, best of 3: 2.96 ms per loop

现在的问题是：你能提出一个更快的代码版本来解决这个问题吗？能够挤出额外的20%将会很棒。

更新于2017年5月

经过相当长的时间，我回到了这个问题，重新运行并扩展了任务和测试。

1. 使用einsum，我已经将代码扩展到P不是一行而是一个矩阵的情况。因此，任务是将O的所有列与P的所有列相关联。

2. 对于如何使用不同的用于科学计算的语言解决相同问题感到好奇，我（在其他人的帮助下）在MATLAB、Julia和R中实现了它。MATLAB和Julia是最快的，并且它们有专门的例程来计算逐列相关性。R也有一个专门的例程，但速度最慢。

3. 在当前版本的numpy（来自Anaconda的1.12.1）中，einsum仍然胜过我使用的专用函数。

所有脚本和时间都可以在https://github.com/ikizhvatov/efficient-columnwise-correlation中找到。

- ilya

---

O和P可以更改吗？此外，这些方法严重依赖于链接到numpy的BLAS库，如果您尚未使用某种优化的BLAS，我建议您使用它。 - Daniel

@Ophion，P 可以更改，O 不能更改。我正在使用链接到Intel MKL的numpy，网址为http://www.lfd.uci.edu/~gohlke/pythonlibs/。 - ilya

现在无法测试，但如果您尝试使用np.einsum('ij,ij->j', DO, DO)代替np.sum(DO ** 2, 0)，以及使用np.dot(DP, DP)代替np.sum(DP ** 2)，可能会有一些边际改进。 - Jaime

Answer 1

我们可以介绍np.einsum来实现这个功能；然而，您的编译方式以及是否使用SSE2可能会有所不同。替换求和操作的额外einsum调用可能看起来是多余的，但是在numpy 1.8之前，numpy ufuncs不使用SSE2，而einsum则使用，我们可以避免一些if语句。

在具有英特尔MKL BLAS的Opteron核心上运行此代码，我得到了奇怪的结果，因为我希望dot调用占据大部分时间。

def newColumnWiseCorrcoef(O, P):
    n = P.size
    DO = O - (np.einsum('ij->j',O) / np.double(n))
    P -= (np.einsum('i->',P) / np.double(n))
    tmp = np.einsum('ij,ij->j',DO,DO)
    tmp *= np.einsum('i,i->',P,P)          #Dot or vdot doesnt really change much.
    return np.dot(P, DO) / np.sqrt(tmp)

O = np.reshape(np.random.rand(100000), (1000,100))
P = np.random.rand(1000)

old = ColumnWiseCorrcoef(O,P)
new = newColumnWiseCorrcoef(O,P)

np.allclose(old,new)
True

%timeit ColumnWiseCorrcoef(O,P)
100 loops, best of 3: 1.52ms per loop

%timeit newColumnWiseCorrcoef(O,P)
1000 loops, best of 3: 518us per loop

在使用英特尔 MKL 的英特尔系统上再次运行此代码，我得到了更合理/预期的结果：

%timeit ColumnWiseCorrcoef(O,P)
1000 loops, best of 3: 524 µs per loop

%timeit newColumnWiseCorrcoef(O,P)
1000 loops, best of 3: 354 µs per loop

再次使用英特尔机器，这次处理的是稍微大一些的东西：

O = np.random.rand(1E5,1E3)
P = np.random.rand(1E5)

%timeit ColumnWiseCorrcoef(O,P)
1 loops, best of 3: 1.33 s per loop

%timeit newColumnWiseCorrcoef(O,P)
1 loops, best of 3: 791 ms per loop

- Daniel

---

tmp = ...; tmp *= ... 这样做会节约时间吗？我看其他人也这么做，但我不明白它比 tmp = ... * ... 有什么优势。 - askewchan

我发现这个解决方案非常有帮助，不知道是否有一种方法可以扩展它以计算回归残差，而不会在速度方面牺牲太多。 - fazistinho_

Answer 2

这是一个效率较低的解决方案，但它非常易读且只有一行代码：

np.einsum("ij,i->j", spst.zscore(A), spst.zscore(B)) / A.shape[0]

如果想要将B作为矩阵而不是向量，并计算对应列的相关性，可以轻松修改为：

np.einsum("ij,ij->j", spst.zscore(A), spst.zscore(B)) / A.shape[0]

根据我在我的iMac上进行的定时测试，配备2.8 GHz四核Intel Core i5处理器，速度大约慢了3倍。

import scipy.stats as spst

def ColumnWiseCorrcoef(O, P):
    n = P.size
    DO = O - (np.sum(O, 0) / np.double(n))
    DP = P - (np.sum(P) / np.double(n))
    return np.dot(DP, DO) / np.sqrt(np.sum(DO ** 2, 0) * np.sum(DP ** 2))

def newColumnWiseCorrcoef(O, P):
    n = P.size
    DO = O - (np.einsum('ij->j',O) / np.double(n))
    P -= (np.einsum('i->',P) / np.double(n))
    tmp = np.einsum('ij,ij->j',DO,DO)
    tmp *= np.einsum('i,i->',P,P)          #Dot or vdot doesnt really change much.
    return np.dot(P, DO) / np.sqrt(tmp)

def OneLineColumnWiseCorrcoef(A, B):
    return np.einsum("ij,i->j", spst.zscore(A), spst.zscore(B)) / A.shape[0]

O = np.reshape(np.random.rand(100000), (1000,100))
P = np.random.rand(1000)

old = ColumnWiseCorrcoef(O,P)
new = OneLineColumnWiseCorrcoef(O, P)
np.allclose(old,new)
True

%timeit ColumnWiseCorrcoef(O,P)
# 325 µs ± 19.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit newColumnWiseCorrcoef(O,P)
# 274 µs ± 713 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit OneLineColumnWiseCorrcoef(O,P)
# 1.02 ms ± 4.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

- calmcc