如何加速Numpy数组中所有通道对的Phase Locking Value计算？

问题分析 #

你的原始代码依赖双重Python循环遍历所有通道对，每次循环单独计算相位差的复指数求和——这种方式完全浪费了Numpy的向量化优势，且Python循环本身效率极低，导致计算耗时过长。结合你128GB的充足内存，可以通过以下几种方式大幅提速：

方法一：向量化矩阵乘法（无NaN场景最优） #

PLV的核心公式可转化为复向量的共轭点积：
$$PLV(a,b) = \frac{|\sum_{t=1}^N e^{i(x_a(t)-x_b(t))}|}{N} = \frac{|\sum_{t=1}^N e^{ix_a(t)} \cdot \overline{e^{ix_b(t)}}|}{N}$$

如果相位数据中无NaN，可直接用矩阵乘法一次性计算所有通道对的结果，底层调用优化后的BLAS/LAPACK库，速度提升数个数量级：

import numpy as np

# 假设x是已加载的256x5000000相位数组
complex_x = np.exp(1j * x)  # 预计算每个通道的复指数向量，形状256x5000000
N_samples = x.shape[1]

# 矩阵乘法计算所有通道对的共轭点积，结果为256x256的复数矩阵
dot_product_matrix = complex_x @ complex_x.conj().T

# 计算最终PLV矩阵
plv_matrix = np.abs(dot_product_matrix) / N_samples

方法二：单循环向量化处理（兼容NaN场景） #

如果相位数据存在NaN，矩阵乘法无法直接处理nansum，可采用单循环+向量化运算的方式，仅保留外层通道循环，内层计算通过Numpy向量化完成，内存占用约10GB（远低于你的128GB上限）：

import numpy as np

complex_x = np.exp(1j * x)
N_samples = x.shape[1]
plv_matrix = np.empty((256, 256), dtype=np.float64)

for a in range(256):
    # 向量化计算当前通道a与所有通道b的复共轭乘积，形状256x5000000
    products = complex_x[a] * complex_x.conj()
    # 对每个通道b的乘积结果做nansum，得到256个复数结果
    sum_products = np.nansum(products, axis=1)
    # 计算PLV值并赋值
    plv_matrix[a] = np.abs(sum_products) / N_samples

如果需要更准确的结果（除以两个通道均非NaN的有效样本数，而非总样本数），可额外计算有效样本数矩阵：

# 生成非NaN掩码
valid_mask = ~np.isnan(x)
# 计算每个通道对的有效样本数（矩阵乘法实现向量化计数）
valid_counts = valid_mask @ valid_mask.T

# 替换PLV计算的分母
plv_matrix[a] = np.abs(sum_products) / valid_counts[a]

方法三：Numba JIT编译（多核加速循环） #

如果上述向量化方式仍有瓶颈，可使用Numba对循环进行JIT编译，开启并行计算利用多核CPU资源：

import numpy as np
from numba import jit, float64, complex128

@jit(float64[:,:](complex128[:,:]), nopython=True, parallel=True)
def compute_plv_numba(complex_x):
    Nchans = complex_x.shape[0]
    N_samples = complex_x.shape[1]
    plv = np.empty((Nchans, Nchans), dtype=np.float64)
    
    # 并行遍历所有通道对
    for a in range(Nchans):
        for b in range(Nchans):
            sum_val = 0j
            for t in range(N_samples):
                val = complex_x[a, t] * np.conj(complex_x[b, t])
                if not np.isnan(val):
                    sum_val += val
            plv[a, b] = np.abs(sum_val) / N_samples
    return plv

# 预计算复指数数组
complex_x = np.exp(1j * x)
# 调用编译后的函数计算PLV
plv_matrix = compute_plv_numba(complex_x)

性能参考 #

• 原始双重Python循环：耗时数小时级别
• 方法一（矩阵乘法）：耗时约几十秒（无NaN场景）
• 方法二（单循环向量化）：耗时约1-2分钟（兼容NaN）
• 方法三（Numba并行）：耗时约数分钟（取决于CPU核心数）

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