如何改进非常低效的numpy代码来计算相关性

时间:2014-01-29 22:33:47

标签: python performance numpy

我编写了以下函数来计算矩阵相对于所选行(由index参数指定)的逐行相关性:

# returns a 1D array of correlation coefficients whose length matches
# the row count of the given np_arr_2d
def ma_correlate_vs_index(np_arr_2d, index):
    def corr_upper(x, y):
        # just take the upper right corner of the correlation matrix
        return numpy.ma.corrcoef(x, y)[0, 1]

    return numpy.ma.apply_along_axis(corr_upper, 1, np_arr_2d, np_arr_2d[index, :])

问题是代码非常非常慢,我不确定如何提高性能。我相信使用apply_along_axis以及corrcoef正在创建2D数组的事实都会导致性能不佳。有没有更直接的计算方法可以提供更好的表现?

如果重要,我使用ma版本的函数来屏蔽数据中找到的一些nan值。此外,我的数据的np_arr_2d形状为(623065, 72)

1 个答案:

答案 0 :(得分:4)

我认为corrcoef中有很多开销是正确的。基本上你只需要每行的点积与索引行,标准化为最大值1.0。

这样的东西会起作用并且速度会快得多:

# Demean
demeaned = np_arr_2d - np_arr_2d.mean(axis=1)[:, None]

# Dot product of each row with index
res = np.ma.dot(demeaned, demeaned[index])

# Norm of each row
row_norms = np.ma.sqrt((demeaned ** 2).sum(axis=1))

# Normalize
res = res / row_norms / row_norms[index]

这比原始代码运行得快得多。我已经使用了蒙面数组方法,因此我认为这将适用于包含NaN的数据。

规范中可能存在微小差异,由ddof中的corrcoef控制,在这种情况下,您可以使用np.ma.std计算row_norms并指定ddof您想。