使用Numba实现快速数组重采样的性能问题与2D版本实现咨询

• 编译开销计入运行时间：@nb.njit 装饰的函数首次运行时需要完成JIT编译，如果你把首次运行的时间计入统计，会包含编译耗时，远高于直接调用Numpy的耗时。
• 现有逻辑以Numpy向量化操作为主：你代码里的repeat、copy、nonzero、interp本身就是Numpy原生的高度优化的C实现，Numba对这类操作的封装反而会引入额外的类型检查、调用开销，不会比直接调用Numpy更快。
• 存在大量不必要的内存操作：你通过先repeat再插值的逻辑实现上采样，会额外生成数倍大小的临时数组（比如5000元素的输入会先生成50000元素的重复数组、再拷贝生成xtmp数组），内存读写开销占比很高，Numba无法优化掉这类冗余的内存操作。

核心优化方向是把冗余的大数组操作替换为直接插值计算，避免不必要的内存开销，这时候Numba的循环优化才能发挥作用，优化后的1D版本示例如下：

import numpy as np
import numba as nb

@nb.njit(cache=True, fastmath=True)
def sync_sampl_1d(time_series, repeat=10):
    n = len(time_series)
    # 直接计算上采样后总长度，不需要生成中间大数组
    out_len = (n - 1) * repeat + 1
    out = np.empty(out_len, dtype=time_series.dtype)
    step = 1.0 / repeat
    for i in range(n-1):
        start_val = time_series[i]
        end_val = time_series[i+1]
        # 填充当前区间的repeat个插值点
        for j in range(repeat):
            out[i*repeat + j] = start_val + (end_val - start_val) * step * j
    # 填充最后一个原始点
    out[-1] = time_series[-1]
    return out

额外优化注意事项：

• 计时前先运行一次函数完成编译预热，避免编译时间干扰测试结果
• 可以根据你的输入数据类型给函数添加静态签名，进一步降低运行时类型检查开销
• 如果不需要严格和原有逻辑完全对齐的边界处理，还可以进一步简化插值计算逻辑，提升速度

• 首先明确上采样维度：如果你的2D数组每行是独立的1D时间序列，仅需要对列方向做上采样，直接循环每行调用优化后的1D上采样函数即可；如果需要同时对行、列两个方向做上采样，按行插值完成后再对列做一次插值即可，对应双线性上采样逻辑。
• 利用Numba并行优化提速：针对4000行的大数组，开启并行循环可以充分利用多核CPU资源，示例代码如下：

@nb.njit(cache=True, fastmath=True, parallel=True)
def sync_sampl_2d(arr, repeat=10, axis=1):
    # axis=1 表示对列方向上采样，axis=0表示对行方向上采样
    if axis == 1:
        n_rows, n_cols = arr.shape
        out_cols = (n_cols - 1) * repeat + 1
        out = np.empty((n_rows, out_cols), dtype=arr.dtype)
        # 并行循环每行
        for i in nb.prange(n_rows):
            out[i] = sync_sampl_1d(arr[i], repeat)
        return out
    elif axis == 0:
        n_rows, n_cols = arr.shape
        out_rows = (n_rows - 1) * repeat + 1
        out = np.empty((out_rows, n_cols), dtype=arr.dtype)
        # 并行循环每列
        for j in nb.prange(n_cols):
            out[:, j] = sync_sampl_1d(arr[:, j], repeat)
        return out
    else:
        raise ValueError("仅支持axis=0或axis=1")

• 如果需要双方向上采样，先调用axis=1做列插值，再调用axis=0做行插值即可，计算效率远高于直接实现2D核插值。

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