使用 Numba 进行较慢的时间序列模拟答案

【问题标题】：Slower time series simulation with Numba使用 Numba 进行较慢的时间序列模拟
【发布时间】：2021-07-21 11:44:53
【问题描述】：

我想在这段代码上使用来自 Numba 的 @njit 装饰器，给定矩阵 A、B、C、D 会从状态空间模型生成样本

x_n = A@x_{n-1} + B@v_n
y_n = C@x_n + D@v_n

@njit
def generate_Y_state_space(N, A, B, C, D):
    """
    Simulate M-dimensional time series given state space model defined by A,B,C,D.
    """
    M = A_sim.shape[0]

    v = np.random.normal(0,1/np.sqrt(2),(M,N)) + 1j*np.random.normal(0,1/np.sqrt(2),(M,N)) # complex gaussian randomly variable
    x = np.zeros((M,N),dtype='c16') # 'c16' is the numba type for complex128
    y = np.zeros((M,N),dtype='c16')

    #initialization
    x[:,0] = v[:,0]
    y[:,0] = C@x[:,0] + D@v[:,0]

    for i in range(1,N):
        x[:,i] = A@x[:,i-1] + B@v[:,i]
        y[:,i] = C@x[:,i] + D@v[:,i]

    return y

但是，没有 njit 装饰器，我得到以下性能（N=1000，M=100）

%timeit generate_Y_state_space(N, A, B, C, D)
27.9 ms ± 728 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

虽然使用了 njit 装饰器，但性能并没有真正提高：

%timeit generate_Y_state_space(N, A, B, C, D)
24.1 ms ± 6.21 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

我想知道矩阵乘法的 Numba 实现是否实际上并不比 Numpy 更好...您知道如何改进此代码吗？

编辑：我认为 Numba 能够提供很好的性能改进，而不是矩阵乘法（因为 Numpy 已经非常快，正如所指出的那样），而是更多的 for 循环（这是必要的，因为整点时间序列的转换是生成一个新数据点作为前一个数据点的转换）。

【问题讨论】：

您的意思是“numba 实现......实际上并不比 numpy 更好”吗？这不会让我感到惊讶。 Numpy 例程已经是高度优化的 C 代码，深入研究汇编程序以进行特定于指令集的优化。 Numba 不会帮助个人操作。
是的，我编辑了我的错误！我知道 numpy 已经非常快了，所以 numba 不会加快很多。同时，需要大量时间的是我认为可以通过使用 Numba 来改进的 for 循环......

标签： python performance numpy time-series numba

【解决方案1】：

使用 Numba 的性能略有下降的一个可能原因是，您至少需要在 @njit 装饰器中使用 fastmath=True 才能与内部使用它的 Numpy 一样快.

另一个原因是@njit 装饰器在运行时编译函数有点慢（通常需要超过 28 毫秒）。您应该注意不要在基准测试中包含此编译时间。您可以指定装饰器中的类型，以便 Numba 可以在第一次调用之前（提前）编译函数。这是一个例子：

@njit('c16[:,::1](int64, c16[:,::1], c16[:,::1], c16[:,::1], c16[:,::1])')

此外，您不需要对数组进行零初始化：x 和 y 可以保持未初始化状态。

最后，您可以使用并行性加快计算速度。这在这里并不简单，因为对x[:,i] 存在时间依赖性。然而，B@v[:,i] 和D@v[:,i] 例如可以并行计算。因此，您可以使用参数parallel=True 和prange 而不是range。

这是一个（未经测试的）示例：

@njit('c16[:,::1](int64, c16[:,::1], c16[:,::1], c16[:,::1], c16[:,::1])', fastmath=True, parallel=True)
def generate_Y_state_space(N, A, B, C, D):
    """
    Simulate M-dimensional time series given state space model defined by A,B,C,D.
    """
    M = A_sim.shape[0]

    v = np.random.normal(0,1/np.sqrt(2),(M,N)) + 1j*np.random.normal(0,1/np.sqrt(2),(M,N)) # complex gaussian randomly variable
    x = np.empty((M,N),dtype='c16') # 'c16' is the numba type for complex128
    y = np.empty((M,N),dtype='c16')

    #initialization
    x[:,0] = v[:,0]
    y[:,0] = C@x[:,0] + D@v[:,0]

    for i in prange(1,N):
        x[:,i] = B@v[:,i]
        y[:,i] = D@v[:,i]

    for i in range(1,N):
        x[:,i] += A@x[:,i-1]
        y[:,i] += C@x[:,i]

    return y

并行性不一定总能让代码更快，但值得在台式机上尝试。

【讨论】：