如何用numpy.vectorize向量化计算两数组对应行的指定区间相关性？

问题背景 #

现有两个二维NumPy数组，需要计算对应行的皮尔逊相关系数，但每行使用的列索引范围由Pandas DataFrame指定：

import numpy as np
import pandas as pd

ts1 = np.array([[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5]])
ts2 = np.array([[6,2,7,4,5],[1,2,3,4,5],[3,2,3,8,5],[2,2,3,1,5],[9,8,7,6,5]])
indices = pd.DataFrame({'in':[0,1,1,2,0],'fin':[4,3,2,3,3]})

已通过Python循环实现需求：

corrIdx = np.zeros((ts1.shape[0],))
for idx in range(ts1.shape[0]):
    corrMatrix = np.corrcoef(ts1[idx,indices['in'][idx]:indices['fin'][idx]],ts2[idx,indices['in'][idx]:indices['fin'][idx]])
    corrIdx[idx] = corrMatrix[0,1].item()

遇到的问题 #

尝试用np.vectorize向量化处理时，两次尝试均报错：

第一次尝试错误 #

代码：

def calcCorr(idx,ts1,ts2,indices):
    corrMatrix = np.corrcoef(ts1[idx,indices['in'][idx]:indices['fin'][idx]],ts2[idx,indices['in'][idx]:indices['fin'][idx]])
    return corrMatrix[0,1].item()

myFunct = np.vectorize(calcCorr)
idx = np.arange(ts1.shape[0])
corrIdx = myFunct(idx,ts1,ts2,indices)

错误信息：

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 2163, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 2241, in _vectorize_call
    ufunc, otypes = self._get_ufunc_and_otypes(func=func, args=args)
  File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 2201, in _get_ufunc_and_otypes
    outputs = func(*inputs)
  File "<stdin>", line 2, in calcCorr
IndexError: invalid index to scalar variable.

第二次尝试错误 #

将索引转为NumPy数组后尝试：

ini = np.array([0,1,1,2,0])
fin = np.array([4,3,2,3,3])

def my_function(a,b,ini,fin):
    result = np.corrcoef(a[ini:fin], b[ini:fin])[0,1]
    return result

vfunc = np.vectorize(my_function, signature='(n),(n),(n),(n)->()')
r = vfunc(ts1,ts2,ini,fin)

错误信息：

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 2163, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 2237, in _vectorize_call
    res = self._vectorize_call_with_signature(func=func, args=args)
  File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 2277, in _vectorize_call_with_signature
    results = func(*(arg[index] for arg in args))
  File "<stdin>", line 2, in my_function
TypeError: only integer scalar arrays can be converted to a scalar index

错误原因 #

1. 第一次尝试：np.vectorize默认会将所有传入的数组参数逐元素拆分，ts1、ts2这类二维数组会被拆成单个标量，导致函数内的ts1[idx, ...]索引操作失效（无法对标量做二维索引）。
2. 第二次尝试：signature参数定义错误，(n),(n),(n),(n)->()表示每个输入都是长度为n的数组，但ini、fin是每行对应一个标量值，并非长度为n的数组，导致切片时用数组作为索引，触发类型错误。
3. 额外注意：np.vectorize本质上还是Python循环的包装，不会带来真正的性能提升，只是语法上看起来像向量化操作。

可行解决方案 #

方案1：用Pandas Apply实现（简单直观） #

将数据合并到DataFrame中，逐行处理：

import numpy as np
import pandas as pd

ts1 = np.array([[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5]])
ts2 = np.array([[6,2,7,4,5],[1,2,3,4,5],[3,2,3,8,5],[2,2,3,1,5],[9,8,7,6,5]])
indices = pd.DataFrame({'in':[0,1,1,2,0],'fin':[4,3,2,3,3]})

# 合并数据到DataFrame
df = indices.copy()
df['ts1'] = list(ts1)
df['ts2'] = list(ts2)

def calc_row_corr(row):
    s1 = row['ts1'][row['in']:row['fin']]
    s2 = row['ts2'][row['in']:row['fin']]
    # 处理切片长度不足2的情况，避免corrcoef报错
    if len(s1) < 2:
        return np.nan
    return np.corrcoef(s1, s2)[0,1]

# 计算相关系数
corrIdx = df.apply(calc_row_corr, axis=1).values
print(corrIdx)

方案2：用Numba加速循环（高性能推荐） #

如果数据量较大，用Numba将循环编译为机器码，大幅提升运行速度：

import numpy as np
import pandas as pd
from numba import jit

ts1 = np.array([[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5]])
ts2 = np.array([[6,2,7,4,5],[1,2,3,4,5],[3,2,3,8,5],[2,2,3,1,5],[9,8,7,6,5]])
indices = pd.DataFrame({'in':[0,1,1,2,0],'fin':[4,3,2,3,3]})
ini = indices['in'].values
fin = indices['fin'].values

@jit(nopython=True)
def compute_corrs(ts1, ts2, ini, fin):
    n_rows = ts1.shape[0]
    corrs = np.zeros(n_rows, dtype=np.float64)
    for i in range(n_rows):
        s1 = ts1[i, ini[i]:fin[i]]
        s2 = ts2[i, ini[i]:fin[i]]
        if len(s1) < 2:
            corrs[i] = np.nan
            continue
        # 手动计算皮尔逊相关系数，减少np.corrcoef的额外开销
        mean1 = np.mean(s1)
        mean2 = np.mean(s2)
        covariance = np.sum((s1 - mean1) * (s2 - mean2))
        std1 = np.sqrt(np.sum((s1 - mean1)**2))
        std2 = np.sqrt(np.sum((s2 - mean2)**2))
        if std1 == 0 or std2 == 0:
            corrs[i] = np.nan
        else:
            corrs[i] = covariance / (std1 * std2)
    return corrs

corrIdx = compute_corrs(ts1, ts2, ini, fin)
print(corrIdx)

方案3：修正np.vectorize用法（不推荐） #

如果坚持要用np.vectorize，需将非循环参数标记为excluded，避免被拆分：

import numpy as np
import pandas as pd

ts1 = np.array([[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5]])
ts2 = np.array([[6,2,7,4,5],[1,2,3,4,5],[3,2,3,8,5],[2,2,3,1,5],[9,8,7,6,5]])
indices = pd.DataFrame({'in':[0,1,1,2,0],'fin':[4,3,2,3,3]})

def calcCorr(idx, ts1, ts2, ini, fin):
    s1 = ts1[idx, ini[idx]:fin[idx]]
    s2 = ts2[idx, ini[idx]:fin[idx]]
    if len(s1) < 2:
        return np.nan
    return np.corrcoef(s1, s2)[0,1].item()

# 排除不需要逐元素处理的参数
myFunct = np.vectorize(calcCorr, excluded=['ts1', 'ts2', 'ini', 'fin'])
idx = np.arange(ts1.shape[0])
corrIdx = myFunct(idx, ts1=ts1, ts2=ts2, ini=indices['in'].values, fin=indices['fin'].values)
print(corrIdx)

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