numpy 2D 数组的快速/矢量化迭代和更新答案

【问题标题】：Fast/vectorized iterating and updating of numpy 2D arraynumpy 2D 数组的快速/矢量化迭代和更新
【发布时间】：2021-10-31 14:58:51
【问题描述】：

这个问题非常贴近我的心，因为我已经做了近 2 年的类似事情，并且一直想知道当行 i 取决于行 i-1 时是否存在修改大型数组/数据帧的矢量化方式，即，当递归听起来是强制性的时。我很想知道是否有聪明的算法或聪明的工具（cython、numba、摆脱冗余操作等） 优化运行时。

问题：
我有 3 个大 numpy 数组：x、y 和 z 形状（100 万乘 500）、（100 万乘 500）和（100 万乘 1）。根据是否为|(x[i] - x[i-1]) * z[i] / y[i]| > thresh，对x 的任何给定行i 中的每个元素进行剪辑/优化。我正在通过以下方式执行此操作，这需要很长时间才能运行我的模拟（尤其是当此步骤重复数千次以调整超参数时）：

t = x.copy()
t[0] = np.clip(t[0] * z[0]/ y[0], -1 * thresh, thresh) * y[0] / z[0]
for i in range(1, t.shape[0]):
    t[i] = np.clip((t[i] - t[i-1]) * z[i] / y[i], -1* thresh, thresh) * y[i] / z[i] + t[i-1]

示例输入：

import numpy as np
import random
x = np.random.rand(1000000, 500)
y = np.random.rand(1000000, 500)
z = np.random.rand(1000000, 1)
thresh = 0.7

编辑：修改以删除 @Mad Physicist 建议的附加和 @Pedro Maia 建议的冗余if-else

【问题讨论】：

为什么你有一个执行 1 英里的 if 语句？如果你只想在循环外移动一次
@PedroMaia 公平点。我现在已经修改了。
不要将附加用于一件事。它每次都重新分配整个数组
考虑使用 cython 或 numba 来完成类似的事情。
@MadPhysicist 谢谢 - 我摆脱了附加。你能告诉我在这种情况下如何使用 numba 和 cython 吗？（我的模拟在 spark 上运行，所以假设 numba/cython 不会在那里引起问题）。非常感谢

标签： python performance vectorization numpy-ndarray

【解决方案1】：

通过使用numba，我得到了 3-4 倍的改进，通过缓存编译的函数又得到了 2 倍（总共 6-8 倍）。

由于我 PC 上的 RAM 较小，我不得不减小 x、y 和 z 的大小。

import numba as nb
import numpy as np

@nb.jit(nopython=True, cache=True)
def bottleneck_func3(t, thresh, y, z):
    t[0] = np.clip(t[0] * z[0] / y[0], -1 * thresh, thresh) * y[0] / z[0]
    for i in range(1, t.shape[0]):
        t[i] = np.clip((t[i] - t[i - 1]) * z[i] / y[i], -1 * thresh, thresh) * y[i] / z[i] + t[i - 1]

x = np.random.rand(300000, 500)
y = np.random.rand(300000, 500)
z = np.random.rand(300000, 1)
thresh = 0.7

t = x.copy()
bottleneck_func3(t, thresh, y, z)

注意 cached function is recompiled each time you modify the file，即使是与缓存功能无关的部分。

【讨论】：