在Numpy中动态折扣累计和

Question

我经常遇到这样一个问题：有两个相同长度的数组，一个包含数值，另一个包含动态衰减因子；并且希望计算每个位置处衰减累积和向量。使用Python循环来表达所需递推关系如下：

import numpy as np

c = np.empty_like(x)
c[0] = x[0]
for i in range(1, len(x)):
    c[i] = c[i-1] * d[i] + x[i]

这段Python代码非常清晰易读，但会显著降低执行速度。为解决此问题，可以通过使用Numba进行JIT编译或Cython预编译来提升性能。如果衰减因子是一个固定的数值（实际情况并非如此），那么可以利用SciPy库中的signals模块，通过lfilter函数实现该功能（参见https://stackoverflow.com/a/47971187/1017986）。

是否存在一种更为“NumPy风格”的方法来表达这一需求，同时不牺牲代码的清晰度和效率呢？

Mark Cidade · Answer

从数学角度讲，这其实是一个具有变系数的非齐次一阶递推关系。

请注意，在您的例子中（即d[1:]中的值），这些系数必须不为零。

这里有一种使用NumPy函数解决该递推关系的方法。需要注意的是，假设d[0]等于1（因为按照你的算法并没有用到它）。在本解决方案中，可以看到c数组中的各个值并不依赖于c数组中的前一个值，即c[i]并不直接依赖于c[i-1]。

g = np.insert(x[1:], 0, 0)
pd = np.cumprod(d)
c = pd * (x[0] + np.cumsum(g / pd))

为了举例说明，我定义了两个函数：一个基于您提供的代码计算递推关系，另一个则利用NumPy实现：

def rec(d, x):
    c = [x[0]]
    for i in range(1, len(x)):
        c.append(d[i] * c[-1] + x[i])
    return np.array(c)

def num(d, x):
    g = np.insert(x[1:], 0, 0)
    pd = np.cumprod(d)
    return pd * (x[0] + np.cumsum(g / pd))

以下是一些使用示例：

>>> rec(np.array([3, 4, 5.2, 6.1]), np.array([2, 3.2, 4, 5.7]))
array([  2.   ,  11.2  ,  62.24 , 385.364])
>>> num(np.array([3, 4, 5.2, 6.1]), np.array([2, 3.2, 4, 5.7]))
array([  2.   ,  11.2  ,  62.24 , 385.364])

>>> rec(np.array([3, 4.3, 5.2]), np.array([6.7, 7.1, 1.2]))
array([  6.7  ,  35.91 , 187.932])
>>> num(np.array([3, 4.3, 5.2]), np.array([6.7, 7.1, 1.2]))
array([  6.7  ,  35.91 , 187.932])

如果您想了解这个答案背后的数学细节（这超出了当前讨论范围），可以参考维基百科上的这篇文章。

编辑

虽然这个解决方案数学上是正确的，但在计算过程中可能会引入数值稳定性问题（特别是当d包含许多接近0的值时，因numpy.cumprod函数的使用而产生此类问题）。

nickf · Answer

这是一个自答，旨在为未来访客提供具体建议。

简而言之：优先选择Numba，避免纯Python实现

通过在对数域中进行计算，可以得到Riccardo Bucco答案的数值稳定性版本。假设d[0] == 1，其实现如下：

def f_numpy_stable(x, d):
    p = np.cumsum(np.log(d))
    return np.exp(p + np.logaddexp.accumulate(np.log(x) - p))

将此版本与纯Python基准实现以及评论和答案（感谢Riccardo Bucco）中的解决方案相比较，在随机输入下得出的结果可以通过np.allclose测试，我认为这足以证明其稳定性。直观上讲，涉及对数域转换的复杂计算应比简单循环慢。让我们验证这一直觉。

结合问题、评论和回答，共有五种解决方案：

纯Python实现（慢速，易读）
使用Numba JIT编译（快速，易读）
使用Cython预编译（快速）
分解为累积乘积和累积求和（纯NumPy实现，可能数值不稳定）
在对数空间中分解计算（纯NumPy实现）

我编写并测试了所有五种方法，使用较大的数组（10k至100M个元素），在我的Intel MacBook pro上的结果是：Numba最快，其次是Cython（比Numba慢最多50%，对于更大的数据集差距更小），然后是Riccardo Bucco的纯NumPy解决方案（大约比Numba慢3倍），出乎意料的是，纯NumPy稳定版紧随其后（比不稳定的版本慢不到10倍），而最慢的是纯Python实现。以下是完整的对比表格（所有时间单位为秒）。

| 数组长度 | Python | 稳定版NumPy | NumPy | Cython | Numba |
|---------|----------|-------------|-----------|-----------|-----------|
| 10,000 | 00.003'840 | 00.000'546 | 00.000'062 | 00.000'030 | 00.000'019 |
| 100,000 | 00.039'600 | 00.005'550 | 00.000'545 | 00.000'296 | 00.000'192 |
| 1,000,000 | 00.401 | 00.056'500 | 00.009'880 | 00.003'860 | 00.002'550 |
| 10,000,000 | 03.850 | 00.590 | 00.092'600 | 00.040'300 | 00.031'900 |
| 100,000,000 | 40.600 | 07.020 | 01.660 | 00.667 | 00.551 |

下面是三种实现所使用的代码，确保在运行timeit之前触发Numba的JIT编译。

以下是对五种不同实现的具体代码：

def f_python(x, d):
    result = np.empty_like(x)
    result[0] = x[0]
    for i in range(1, x.shape[0]):
        result[i] = result[i-1] * d[i] + x[i]
    return result

@numba.jit
def f_numba(x, d):
    result = np.empty_like(x)
    result[0] = x[0]
    for i in range(1, x.shape[0]):
        result[i] = result[i-1] * d[i] + x[i]
    return result

def f_numpy(x, d):
    result = np.cumprod(d)
    return result * np.cumsum(x / result)

def f_numpy_stable(x, d):
    p = np.cumsum(np.log(d))
    return np.exp(p + np.logaddexp.accumulate(np.log(x) - p))

Cython实现则在另一个笔记本单元格中单独编译：

%%cython
import numpy as np
cimport numpy as np

cpdef np.ndarray[np.float64_t, ndim=1] f_cython(np.ndarray[np.float64_t, ndim=1] x, np.ndarray[np.float64_t, ndim=1] d):
    cdef:
        int i = 0
        int N = x.shape[0]
        np.ndarray[np.float64_t, ndim=1] result = np.empty_like(x)
    result[0] = x[0]
    for i in range(1, N):
        result[i] = result[i-1] * d[i] + x[i]
    return result

我对Numba表现优于Cython感到有些惊讶。考虑到将递推关系分解为独立的累积乘积和累积求和可能会导致数值不稳定或速度相对较慢（尽管仍然比纯Python快），我的建议是尽量避免这种做法（特别是在大数据集情况下），转而采用Cython或Numba。鉴于仅需导入Numba并添加装饰器就能获得更快的速度，因此对我来说，选择Numba是一个显而易见的选择。

非常乐意听取改进Cython实现的所有建议！ :-)