高维时间序列距离相关自相关计算的内存与并行优化问题

问题背景 #

现有两个数组：

• 坐标数组 X，形状为 X.shape = (100000, 1000)
• 属性数组 P，形状为 P.shape = (100000, 1000)

需要计算P的距离相关自相关函数：

• 基础自相关计算逻辑：

  Autocorrelation = [np.mean(P[tau:]*P[:-tau], axis=(0, 1)) for tau in range(P.shape[0])]
  

• 距离相关自相关：对每个偏移量tau，遍历所有时间点t，找到满足np.abs(X[t, i]-X[t+tau, j]) == d的索引i,j，收集对应P[t,i]*P[t+tau,j]的值，最终按tau和d计算平均值。

初始实现与并行异常 #

针对单个tau编写了如下函数，并用Joblib的Parallel进行并行计算：

import numpy as np
from joblib import Parallel, delayed  

X = np.ones([100000, 1000])
P = np.ones([100000, 1000])

def Myfunc(tau):
    AutocorrelationP = []
    for t in range(X.shape[0]-tau):
        d = np.round(np.sqrt((X[t][:, None]-X[t+tau])**2))
        Pd = []
        for i in np.arange(1, 11):   
            Indices = np.where(d==i)
            Pd.append(np.mean(P[t, Indices[0]]*P[t+tau, Indices[1]]))
        AutocorrelationP.append(Pd)
    return np.mean(np.array(AutocorrelationP), axis=0)   

• 当使用10个CPU核心时，脚本运行正常，htop显示所有CPU负载达100%（绿色，计算密集状态）。
• 但使用节点全部64个CPU时，并行开销大幅增加：任务数≤10时符合预期，超过10后运行时间随任务数近似线性增长。

测试不同数组维度后发现：

• 仅当X.shape[1]较大时出现该问题，X.shape[1]≈100时开销消失；
• 极端情况X.shape[1]=5000时，所有CPU执行时变为红色（内存/IO等待状态），判断内存是核心问题，X.shape[0]也有较小但显著的影响。

尝试的解决方案及问题 #

Numba显式循环实现 #

为避免构建大型距离矩阵，用Numba装饰的三层循环替代Numpy广播：

from numba import jit

@jit(nopython=True)
def myFunc(tau):
    PAutocorrelation = np.zeros(11)
    for t in range(X.shape[0]-tau):  
        for i in range(X.shape[1]):
             for j in range(X.shape[1]):
                  dist = int(np.rint(np.abs(X[t, i]-X[t+tau,j])))
                  PAutocorrelation[dist] += P[t, i]*P[t+tau, j]  
    return PAutocorrelation   

• 小数组测试中，该方法的并行开销表现更好；
• 但使用原数组维度时内核崩溃，疑惑：明明没有构建额外大型矩阵，为何仍存在内存问题？

Fast Histogram优化方案 #

尝试用fast_histogram统计距离与权重的关系：

from fast_histogram import histogram1d

def myFunc(tau):
    PAutocorrelation = []
    for t in range(100):
         PAutocorrelation.append(
             histogram1d(np.abs(X[t][:, None]-X[t+tau]).flatten(),
                         bins=10, range=[0., 10],
                         weights=(P[t][:,None]*P[t+tau]).flatten())
         )    
    return PAutocorrelation  

该方案仍存在与第一种方法相同的内存问题。

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