具有运行值的矢量化循环取决于先前的值（+ if 语句）

Question

我习惯于在 Python 中编写向量化语句和列表理解, 以及 if 语句。示意图如下所示：

def my_loop(x, a=0.5, b=0.9):
  out = np.copy(x)
  prev_val = 0 
  for i in np.arange(x.shape[0]):

      if x[i] < prev_val:
          new_val = (1-a)*x[i] + a*prev_val
      else:
          new_val = (1-b)*x[i] + b*prev_val

      out[i] = new_val

      prev_val = new_val

  return out

我一直无法弄清楚如何对其进行矢量化（例如，通过使用某种累加器），所以我想问：有没有办法使它更多 Pythonic/faster？

我以前看过关于在有 if 语句时进行矢量化的帖子——通常通过 np.where() 解决——但没有一个 "running" 值取决于其先前状态...所以我还没有发现任何重复的问题（ isn't about vectorization in the usual sense, 是关于 'previous value' 但指的是列表索引）。

到目前为止，我已经尝试了 np.vectorize 和 numba 的 @jit，它们的速度运行稍快一些，但都没有达到我希望的速度。有什么我想念的吗？（也许是 map()？）谢谢。

（是的，在 a=b 的情况下，这变得很容易！）

Answer 1

我意识到通过删除虚拟变量，可以将此代码放入一种形式，其中 numba 和 @autojit 可以发挥它们的魔力并使其成为 "fast":

from numba import jit, autojit

@autojit
def my_loop4(x, a=0.5, b=0.9):
  out = np.zeros(x.shape[0],dtype=x.dtype)
  for i in np.arange(x.shape[0]):
      if x[i] < out[i-1]:
          out[i] = (1-a)*x[i] + a*out[i-1]
      else:
          out[i] = (1-b)*x[i] + b*out[i-1]
  return out

没有@autojit，这仍然很慢。但是有了它，...问题就解决了。因此，删除不必要的变量 and 添加 @autojit 是成功的方法。

Answer 2

在 nopython 模式下 JITing 更快。引自 numba 文档：

Numba has two compilation modes: nopython mode and object mode. The former produces much faster code, but has limitations that can force Numba to fall back to the latter. To prevent Numba from falling back, and instead raise an error, pass nopython=True.

@nb.njit(cache=True)
def my_loop5(x, a=0.5, b=0.9):
  out = np.zeros(x.shape[0],dtype=x.dtype)
  
  for i in range(x.shape[0]):
      if x[i] < out[i-1]:
          out[i] = (1-a) * x[i] + a * out[i-1]
      else:
          out[i] = (1-b) * x[i] + b * out[i-1]
  return out

因此：

x = np.random.uniform(low=-5.0, high=5.0, size=(1000000,))

时间是：

my_loop4 : 0.235s

my_loop5 : 0.193s

HTH.

具有 运行 值的矢量化循环取决于先前的值（+ if 语句）

Vectorizing loop with running value that depends on previous value (+ if statement)

python

numba

具有运行值的矢量化循环取决于先前的值（+ if 语句）