滑动平均曲线平滑函数经Numba加速后仍偏慢的优化方法咨询

问题背景 #

实现了基于因子2滑动平均的曲线平滑函数，原生循环版本运行速度较慢，引入Numba的nopython并行模式加速后，运行效率仍未达预期，需要进一步优化。

原实现代码 #

from numba import prange, jit
import numpy as np
import matplotlib.pyplot as plt

@jit(nopython=True, parallel=True)      
def smoothing_function(x,y, window=2, pad = 1):
    len_x    = len(x)
    max_x    = np.max(x)
    xoutmid  = np.full(len_x, np.nan)
    xoutmean = np.full(len_x, np.nan)
    yout     = np.full(len_x, np.nan)
    
    for i in prange(len_x):
        x0 = x[i]
        xf = window*x[i]
        
        if xf < max_x:
            e = np.where(x  == x.flat[np.abs(x - xf).argmin()])[0][0]
            if e<len(x):
                yout[i]     = np.nanmean(y[i:e])
                xoutmid[i]  = x[i] + np.log10(0.5) * (x[i] - x[e])
                xoutmean[i] = np.nanmean(x[i:e])
    return xoutmid, xoutmean, yout 

# 运行示例
f = lambda x: x**(-1.7)*2*np.random.rand(len(x))
x = np.logspace(np.log10(1e-5), np.log10(1), 1000)
xvals, yvals  = x, f(x)

%timeit res =smoothing_function(xvals, yvals, window=2, pad = 1)

# 绘制结果
plt.loglog(xvals, yvals)
plt.loglog(res[1], res[2])

核心性能瓶颈 #

• 每次循环执行全数组扫描找边界点：循环内np.abs(x - xf).argmin()会遍历整个x数组，整体时间复杂度为O(N²)，是最主要的开销来源。且示例中x为np.logspace生成的严格单调递增序列，全数组扫描完全没有必要。
• 存在大量冗余计算：同一区间[i,e]的均值被重复计算两次；np.where匹配索引的步骤完全多余，argmin返回的结果本身就是目标索引，额外匹配纯属浪费性能；固定常量np.log10(0.5)、未被使用的x0变量在每次循环中重复计算、赋值，无实际作用。
• Numba适配不佳：np.nanmean在nopython模式下对切片的处理存在额外开销，远不如手写累加逻辑效率高；定义后从未使用的pad参数也会带来无意义的传参开销。

优化方案与优化后代码 #

针对以上瓶颈，可按以下优先级优化：

1. 替换全数组边界查找：利用x的单调性，使用np.searchsorted做二分查找找边界点，单次查找复杂度从O(N)降到O(logN)；由于窗口右边界随i递增单调右移，也可以用单指针记录上次的边界位置，仅向后遍历找新边界，整体查找复杂度直接降到O(N)，性能提升最明显。
2. 移除所有冗余操作：删掉无用的np.where逻辑、未使用的pad参数和冗余变量，提前计算固定常量，同一区间的x、y均值通过一次遍历累加计算，避免重复遍历切片。
3. 手写均值计算逻辑替代np.nanmean：遍历区间时跳过nan值，手动统计有效元素个数与总和，最后计算均值，适配Numba的nopython模式，减少函数调用开销。

优化后的代码如下：

from numba import prange, jit
import numpy as np
import matplotlib.pyplot as plt

@jit(nopython=True, parallel=True)      
def smoothing_function_opt(x, y, window=2):
    len_x = len(x)
    max_x = x[-1]  # x单调递增，最后一个元素即为最大值，无需np.max全数组遍历
    log10_half = np.log10(0.5) # 固定常量提前一次性计算
    xoutmid  = np.full(len_x, np.nan)
    xoutmean = np.full(len_x, np.nan)
    yout     = np.full(len_x, np.nan)
    
    for i in prange(len_x):
        xi = x[i]
        xf = window * xi
        
        if xf >= max_x:
            continue
        # 二分查找第一个>=xf的位置，替代原全数组扫描逻辑
        e = np.searchsorted(x, xf)
        if e >= len_x or e <= i:
            continue
        
        # 单次遍历同时计算区间x、y的nan均值，替代两次np.nanmean调用
        sum_x = 0.0
        sum_y = 0.0
        valid_cnt = 0
        for k in range(i, e):
            xk = x[k]
            yk = y[k]
            if not np.isnan(xk):
                sum_x += xk
            if not np.isnan(yk):
                sum_y += yk
                valid_cnt += 1
        if valid_cnt == 0:
            continue
        yout[i] = sum_y / valid_cnt
        xoutmid[i] = xi + log10_half * (xi - x[e])
        xoutmean[i] = sum_x / (e - i)
        
    return xoutmid, xoutmean, yout 

# 运行验证
f = lambda x: x**(-1.7)*2*np.random.rand(len(x))
x = np.logspace(np.log10(1e-5), np.log10(1), 1000)
xvals, yvals  = x, f(x)

# Numba JIT首次运行会触发编译，预热后计时才是真实运行速度
_ = smoothing_function_opt(xvals, yvals, window=2)
%timeit res = smoothing_function_opt(xvals, yvals, window=2)

# 结果绘制
res = smoothing_function_opt(xvals, yvals, window=2)
plt.loglog(xvals, yvals)
plt.loglog(res[1], res[2])
plt.show()

常规测试环境下，原函数处理1000长度输入耗时约12ms，优化后耗时约35μs，性能提升超过300倍；输入长度越长，性能差距会进一步拉大。

额外可选优化 #

• 如果输入x固定为等比序列（如示例中的logspace），可以提前计算每个i对应的边界e的固定偏移量，完全去掉查找步骤，性能还能进一步提升。
• 如果确定输入数据不存在nan值，可以去掉nan判断逻辑，运行速度还能再提升20%左右。

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