Python 是否可以使用矢量化加速 numpy 中的这种平均自相关计算

我假设您主要关心较小的运行时间，而不太关心它是如何实现的.以下方法通过改进相关系数的计算，实现了3倍以上的加速.

更具体地说，它将对np.corrcoef的调用替换为对自定义函数fast_corrcoef的调用.本质上，fast_corrcoef遵循numpy方法来计算相关系数，只是它不执行大量冗余计算.它主要利用了这样一个事实:对于调用np.corrcoef的每个调用(调用np.cov来计算协方差矩阵，需要计算S[:-lag]和S[lag:]系列的平均值)，都会进行大量不必要的计算:而不是像np.cov中那样单独计算平均值，fast_corrcoef通过首先计算S[lag:-lag]的和，然后使用中间和计算两个系列的平均值来计算它们.对于大小为n的时间序列，这只产生(n - 2 * lag) + 2 * lag + 2 = n + 2个加法，而不是2n个加法，即在A.shape = (100, 3000)和lag = 5个加法总计(2 * (100 - 5) - (100 - 5 + 2)) * 3000 = 279000个加法的情况下，基本上节省了所有加法中的50%个用于计算平均值.

此外，由于协方差矩阵是对称的，因此计算协方差矩阵时不需要全矩阵乘法，因此，在np.cov的情况下对np.dot的调用也被替换为只进行实际需要的树向量点积的表达式，即只计算一次协方差，而不是两次.

import numpy as np
import time


def fast_corrcoef(X, lag):
    n = X.shape[1]
    mid_sum = X[0, lag:].sum()
    X -= np.array([(mid_sum + X[0, :lag].sum()) / n, (mid_sum + X[1, -lag:].sum()) / n])[:, None]
    return (X[0,:] * X[1,:]).sum() / (np.sqrt((X[0,:] ** 2).sum()) * np.sqrt((X[1,:] ** 2).sum()))
    
A = np.random.rand(100,3000)

lag = 5

def orig(A):
    return np.nanmean(np.array([np.corrcoef(np.array([A[:-5,i],A[5:,i]]))[1,0] for i in range(A.shape[1])]),axis=0)

def faster(A, lag):
    corrcoefs = np.empty([A.shape[1]])
    for i in range(A.shape[1]):
        corrcoefs[i] = fast_corrcoef(np.array([A[:-lag, i], A[lag:, i]]), lag)
    return np.nanmean(corrcoefs, axis=0)

start_time = time.time()
res_orig = orig(A)
runtime_orig = time.time() - start_time
print(f'orig: runtime: {runtime_orig:.4f}s; result: {res_orig}')


start_time = time.time()
res_faster = faster(A, lag)
runtime_faster = time.time() - start_time
print(f'fast: runtime: {runtime_faster:.4f}s; result: {res_faster}')

assert abs(res_orig - res_faster) < 1e-16

speedup = runtime_orig / runtime_faster

print(f'speedup: {speedup:.4f}x')

这段代码打印了以下内容(当我在Google Colab中运行它时):

orig: runtime: 0.4214s; result: -0.00961872402216293
fast: runtime: 0.1138s; result: -0.009618724022162932
speedup: 3.7030x