如何用向量化或其他方法优化Numpy数组峰值检测以替代循环？

我有一个1080万行的NumPy数组，第一列为时间数据，其余10列是信号值。目标是对所有信号列执行峰值检测，获取峰值及其对应的时间。

现有循环实现代码可正常运行，但耗时过长：

import numpy as np
from peakutils import indexes

arr = stress_data.to_numpy()
times_all = arr[:,0]
peaks_all = []
width = 125

for i in range(1,np.shape(arr)[1]):
    x = arr[:,i]
    xf = x - np.mean(x)    
    threshold = 0.1*np.average(xf) / np.max(xf)
   
    # Find x-coordinates of peaks in signal
    peaks = indexes(xf, thres = threshold, min_dist = width)
        
    sg = [times_all[peaks], xf[peaks]]
    peaks_all.append(sg)

尝试用np.apply_along_axis优化，结果运行速度反而更慢：

def process_column(x):
        xf = x - np.mean(x)
        threshold = 0.1 * np.average(xf) / np.min(xf)
        valleys = indexes(xf, thres=threshold, min_dist=width)
        return [times_all[valleys], xf[valleys]]

valleys_all = np.apply_along_axis(process_column, axis=0, arr=data_all.T)

数据示例：

1. 批量预计算统计量，减少循环内重复计算 #

将循环内单列的均值、最大值计算改为批量处理，大幅降低重复计算开销：

import numpy as np
from peakutils import indexes

arr = stress_data.to_numpy()
times_all = arr[:, 0]
signal_cols = arr[:, 1:]  # 提取所有信号列
width = 125

# 批量计算所有信号列的均值
col_means = np.mean(signal_cols, axis=0)
# 批量生成去均值后的信号矩阵
xf_matrix = signal_cols - col_means[np.newaxis, :]
# 批量计算每列的最大值和均值（保留原逻辑）
col_max = np.max(xf_matrix, axis=0)
col_avg = np.average(xf_matrix, axis=0)
thresholds = 0.1 * col_avg / col_max

peaks_all = []
for idx in range(signal_cols.shape[1]):
    xf = xf_matrix[:, idx]
    threshold = thresholds[idx]
    peaks = indexes(xf, thres=threshold, min_dist=width)
    peaks_all.append([times_all[peaks], xf[peaks]])

2. 替换peakutils.indexes为更高效的scipy.signal.find_peaks #

peakutils的峰值检测实现效率低于scipy的优化版本，替换后可显著提升单列处理速度：

import numpy as np
from scipy.signal import find_peaks

arr = stress_data.to_numpy()
times_all = arr[:, 0]
signal_cols = arr[:, 1:]
width = 125

col_means = np.mean(signal_cols, axis=0)
xf_matrix = signal_cols - col_means[np.newaxis, :]
col_max = np.max(xf_matrix, axis=0)
col_avg = np.average(xf_matrix, axis=0)
thresholds = 0.1 * col_avg / col_max

peaks_all = []
for idx in range(signal_cols.shape[1]):
    xf = xf_matrix[:, idx]
    # find_peaks的height对应阈值，distance对应min_dist
    peaks, _ = find_peaks(xf, height=thresholds[idx], distance=width)
    peaks_all.append([times_all[peaks], xf[peaks]])

3. 并行处理多列信号 #

利用多进程并行处理独立的信号列，充分发挥CPU多核性能：

import numpy as np
from scipy.signal import find_peaks
from concurrent.futures import ProcessPoolExecutor

def process_single_col(args):
    xf, threshold, times_all, width = args
    peaks, _ = find_peaks(xf, height=threshold, distance=width)
    return [times_all[peaks], xf[peaks]]

arr = stress_data.to_numpy()
times_all = arr[:, 0]
signal_cols = arr[:, 1:]
width = 125

col_means = np.mean(signal_cols, axis=0)
xf_matrix = signal_cols - col_means[np.newaxis, :]
col_max = np.max(xf_matrix, axis=0)
col_avg = np.average(xf_matrix, axis=0)
thresholds = 0.1 * col_avg / col_max

# 构造并行任务参数
tasks = [(xf_matrix[:, idx], thresholds[idx], times_all, width) for idx in range(signal_cols.shape[1])]

# 启动多进程处理
with ProcessPoolExecutor() as executor:
    peaks_all = list(executor.map(process_single_col, tasks))

4. 原代码潜在问题提示 #

• 去均值后的信号xf均值恒为0，导致np.average(xf)为0，计算出的threshold为0，可能不符合峰值检测预期，建议检查阈值逻辑。
• 若需检测谷值，可使用find_peaks(-xf)或调整scipy.signal.find_peaks的参数。

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