如何加速基于1/3倍频程窗口的数据平滑函数？

原函数用于基于1/3倍频程窗口实现数据平滑，但面对百万级数据时运行极慢：

import numpy as np

def getRollingMean(data, freq, fLim):
    meanData = []
    for fl in fLim:
        meanData.append(np.mean(data[(freq > fl[0]) & (freq < fl[1])]))
    return np.array(meanData)

参数说明 #

• data: 待平滑的一维数组
• freq: 与data长度一致的频率轴数组
• fLim: 二维数组，每行对应一个1/3倍频程窗口的[下限, 上限]

示例 #

data = [1,5,3,0]
freq = [100,150,225,337.5]
fLim = [(50, 200),
        (75,300),
        (112, 450),
        (168,674)]

# 调用后结果
mean = [3,3,2.6,3]

优化方案 #

1. 前缀和+二分查找（适用于freq单调场景）

频率轴通常是单调递增的，利用这个特性可以用二分查找快速定位窗口边界，结合前缀和实现O(1)区间均值计算，彻底避免Python循环：

import numpy as np

def getRollingMean_fast(data, freq, fLim):
    data = np.asarray(data)
    freq = np.asarray(freq)
    fLim = np.asarray(fLim)
    
    # 计算前缀和与前缀计数
    prefix_sum = np.concatenate([[0], np.cumsum(data)])
    prefix_count = np.concatenate([[0], np.cumsum(np.ones_like(data))])
    
    # 二分查找定位每个窗口的左右边界索引
    left_idx = np.searchsorted(freq, fLim[:, 0], side='right')
    right_idx = np.searchsorted(freq, fLim[:, 1], side='left')
    
    # 计算区间均值，处理空窗口避免除0
    interval_sum = prefix_sum[right_idx] - prefix_sum[left_idx]
    interval_count = prefix_count[right_idx] - prefix_count[left_idx]
    return np.where(interval_count > 0, interval_sum / interval_count, np.nan)

该方法时间复杂度为O(M logN)（M为窗口数，N为数据长度），比原O(M*N)的效率提升几个数量级。

2. 向量化广播（无需freq单调）

如果频率轴不单调，可通过广播机制一次性生成所有窗口的掩码矩阵，用numpy内置的向量化运算替代Python循环：

import numpy as np

def getRollingMean_vectorized(data, freq, fLim):
    data = np.asarray(data)
    freq = np.asarray(freq)
    fLim = np.asarray(fLim)
    
    # 广播生成(M, N)的掩码矩阵，标记每个数据点属于哪个窗口
    mask = (freq[np.newaxis, :] > fLim[:, 0, np.newaxis]) & (freq[np.newaxis, :] < fLim[:, 1, np.newaxis])
    
    # 计算每个窗口的均值，自动忽略空窗口
    sum_vals = np.sum(data * mask, axis=1)
    count_vals = np.sum(mask, axis=1)
    return np.where(count_vals > 0, sum_vals / count_vals, np.nan)

注意：当窗口数M和数据量N都很大时，掩码矩阵会占用较多内存，适合窗口数量较少的场景。

3. Numba JIT编译（兼容原逻辑）

如果无法做向量化改造，可用Numba将原函数编译为机器码，直接加速Python循环：

import numpy as np
from numba import jit

@jit(nopython=True)
def getRollingMean_numba(data, freq, fLim):
    meanData = np.zeros(len(fLim))
    for i in range(len(fLim)):
        fl_low, fl_high = fLim[i]
        total = 0.0
        count = 0
        for j in range(len(data)):
            if fl_low < freq[j] < fl_high:
                total += data[j]
                count += 1
        meanData[i] = total / count if count > 0 else np.nan
    return meanData

Numba会跳过Python解释器直接执行编译后的机器码，循环效率可提升10~100倍。

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