元素量级差异大的NumPy数组滑动求和精度优化问询

问题背景 #

需要计算长度为l的数组a的滑动求和：给定步长n，生成长度为l-n+1的数组b，每个元素是a中连续n个元素的和。

原实现方法在元素量级相近时正常，但当数组包含量级差异极大的元素（如1e8和1e-8）时，np.cumsum的累加误差会导致结果精度严重丢失：

def run_sum(a, n):     
    a_cumsum = np.r_[0, np.cumsum(a)]     
    b = a_cumsum[n:] - a_cumsum[:-n]     
    return b

由于数组长度达1e8，无法用循环实现，尝试过以下索引方式但无效：

ind = np.arange(a.size - n + 1)
b[ind] = np.sum(a[ind:ind+n])

精度损失示例 #

a = np.array([-9999.999966666666, -5.7847991744735696e-08, -2.983181169098924e-08, -2.093995193995538e-08, -1.6894298480290635e-08, -1.4869270858640853e-08, -1.3961300790320832e-08, -1.3834651535607362e-08, -1.4395718046705324e-08, -1.570696320519037e-08])

run_sum(a, 3) = 
np.array([-9999.999966754345, -1.0861913324333727e-07, -6.766640581190586e-08, -5.2703398978337646e-08, -4.5723936636932194e-08, -4.2664396460168064e-08, -4.219145921524614e-08, -4.393768904265016e-08])

a[:-2] + a[1:-1] + a[2:] = 
np.array([-9999.999966754345, -1.0861975537568031e-07, -6.766606211123524e-08, -5.2703521278886865e-08, -4.572487012925232e-08, -4.266522318456905e-08, -4.2191670372633516e-08, -4.393733278750305e-08])

上述精度差异对业务至关重要。

改进方案：按量级拆分计算 #

提出run_sum_new方法，通过按元素量级拆分数组，分别计算累加和后再合并，以最小化误差：

def run_sum_new(a, n, max_exp_sep=3, exp_sep_lim=4):
    """
    计算数组a中连续n个元素的滑动求和。当数组元素量级差异较大时，将元素拆分到不同数组中计算，以降低误差。

    参数：
    a - 输入数组
    n - 滑动窗口大小
    max_exp_sep - 正负元素各自最多拆分的数组数量，总数组数量最多为2*max_exp_sep+1
    exp_sep_lim - 元素量级差的最小值，超过该值则拆分到不同数组

    返回：
    run_sum - 滑动求和结果数组
    """

    # 计算正负元素的最大量级和所有元素的最小量级
    a_pos = a[a>0]
    a_neg = a[a<0]
    if a_pos.size == 0:
        a_pos_max = np.nan
    else:
        a_pos_max = np.max(a_pos)
    if a_neg.size == 0:
        a_neg_max = np.nan
    else:
        a_neg_max = np.max(np.abs(a_neg))
    a_min = np.min(np.abs(a))
    exp_pos_max = np.ceil(np.log10(a_pos_max))
    exp_neg_max = np.ceil(np.log10(a_neg_max))
    exp_min = np.floor(np.log10(a_min))
    del(a_pos)
    del(a_neg)

    # 生成拆分阈值列表
    d_exp_pos = exp_pos_max - exp_min
    if np.isnan(d_exp_pos):
        exp_pos_sep_list = []
    elif d_exp_pos <= exp_sep_lim * (max_exp_sep + 1):
        exp_pos_sep_list = [10**(exp_min + exp_sep_lim * i) for i in range(1, max_exp_sep+1) if 10**(exp_min + exp_sep_lim * i) < a_pos_max]
    else:
        new_exp_sep_lim = np.ceil(d_exp_pos / (max_exp_sep + 1))
        print(f"警告：数组正元素的量级跨度{d_exp_pos}超过给定max_exp_sep={max_exp_sep}和exp_sep_lim={exp_sep_lim}的限制。")
        print(f"已采用新的exp_sep_lim={new_exp_sep_lim}。")
        exp_pos_sep_list = [10**(exp_min + new_exp_sep_lim * i) for i in range(1, max_exp_sep+1)]
    d_exp_neg = exp_neg_max - exp_min
    if np.isnan(d_exp_neg):
        exp_neg_sep_list = []
    elif d_exp_neg <= exp_sep_lim * (max_exp_sep + 1):
        exp_neg_sep_list = [-10**(exp_min + exp_sep_lim * i) for i in range(1, max_exp_sep+1) if 10**(exp_min + exp_sep_lim * i) < a_neg_max]
    else:
        new_exp_sep_lim = np.ceil(d_exp_neg / (max_exp_sep + 1))
        print(f"警告：数组负元素的量级跨度{d_exp_neg}超过给定max_exp_sep={max_exp_sep}和exp_sep_lim={exp_sep_lim}的限制。")
        print(f"已采用新的exp_sep_lim={new_exp_sep_lim}。")
        exp_neg_sep_list = [10**(exp_min + new_exp_sep_lim * i) for i in range(1, max_exp_sep+1)]
    exp_neg_sep_list.reverse()
    exp_sep_list = [-np.inf, ] + exp_neg_sep_list + exp_pos_sep_list + [np.inf,]

    # 拆分原数组
    a_sep_list = [np.where((i <= a) & (a < j), a, 0) for (i, j) in zip(exp_sep_list[:-1], exp_sep_list[1:])]
    a_sep_arr = np.array(a_sep_list)

    # 计算滑动求和
    a_sep_cumsum = np.cumsum(a_sep_arr, axis=-1)
    a_sep_cumsum = np.hstack(([[0]]*a_sep_cumsum.shape[0], a_sep_cumsum))
    run_sum = np.sum(a_sep_cumsum[:,n:] - a_sep_cumsum[:,:-n], axis=0)

    return run_sum

测试结果 #

run_sum_new(a, 3) = np.array([-9999.999966754345, -1.0861975537568031e-07, -6.766606211123526e-08, -5.270352127888686e-08, -4.5724870129252317e-08, -4.266522318456904e-08, -4.2191670372633516e-08, -4.393733278750304e-08])

a[:-2] + a[1:-1] + a[2:] = 
np.array([-9999.999966754345, -1.0861975537568031e-07, -6.766606211123524e-08, -5.2703521278886865e-08, -4.572487012925232e-08, -4.266522318456905e-08, -4.2191670372633516e-08, -4.393733278750305e-08])

timeit.timeit('run_sum_new(a, 3)', number=1000, globals=globals()) = 0.12470354699999997

现需确认该方案的有效性，或寻求更优的高精度滑动求和实现方式。

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