如何高效计算Numpy数组的滚动几何均值？

滚动算术均值可以用Numpy的convolve快速实现，但滚动几何均值需要换个思路——利用对数将乘积运算转换为求和运算，再结合Numpy的高效累积求和功能来实现，具体步骤如下：

核心思路 #

几何均值的计算公式是窗口内元素乘积的1/窗口大小次方，而对数运算能把乘积转为求和：ln(a*b*c) = ln(a)+ln(b)+ln(c)，我们可以先对数组取对数，计算滚动窗口的对数和，再取指数还原得到几何均值。

代码实现 #

import numpy as np

def rolling_geometric_mean(a, window_size):
    # 对数组取自然对数，将乘积操作转为求和
    log_vals = np.log(a)
    # 计算对数数组的累积和，用于快速计算窗口内的和
    cum_log = np.cumsum(log_vals)
    
    # 生成每个位置对应的窗口大小：前window_size-1个位置窗口从1递增到window_size-1，后续固定为window_size
    window_lengths = np.concatenate([
        np.arange(1, window_size),
        np.full(len(a) - window_size + 1, window_size)
    ])
    
    # 计算每个窗口的对数和
    rolling_log_sum = []
    for idx in range(len(a)):
        win_len = window_lengths[idx]
        if idx >= win_len:
            current_sum = cum_log[idx] - cum_log[idx - win_len]
        else:
            current_sum = cum_log[idx]
        rolling_log_sum.append(current_sum)
    
    # 转换为几何均值：指数(对数和 / 窗口大小)
    geometric_means = np.exp(np.array(rolling_log_sum) / window_lengths)
    # 按需保留小数位数
    return np.round(geometric_means, 2)

# 测试示例
test_array = np.array([0.5, 2.0, 4.0])
window = 2
print(rolling_geometric_mean(test_array, window))  # 输出: [0.5  1.   2.83 4.  ]

注意事项 #

• 输入数组的元素必须全为正数，因为对数函数的定义域不包含非正数；
• 利用Numpy的cumsum和向量化运算，比纯Python循环遍历窗口的效率高得多，适合处理大规模数组。

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