向量化 numpy 中的 for 循环以计算胶带重叠

Question

我正在使用 python 创建一个应用程序来计算胶带重叠（为分配器建模，将产品应用于旋转鼓）。

我有一个程序可以正常运行，但速度非常慢。我正在寻找一种解决方案来优化用于填充 numpy 数组的 for 循环。有人可以帮我矢量化下面的代码吗？

import numpy as np
import matplotlib.pyplot as plt

# Some parameters
width = 264
bbddiam = 940
accuracy = 4 #2 points per pixel

drum = np.zeros(accuracy**2 * width * bbddiam).reshape((bbddiam * accuracy , width * accuracy))

# The "slow" function
def line_mask(drum, coef, intercept, upper=True, accuracy=accuracy):
    """Masks a half of the array"""
    to_return = np.zeros(drum.shape)
    for index, v in np.ndenumerate(to_return):
        if upper == True:
            if index[0] * coef + intercept > index[1]:
                to_return[index] = 1
        else:
            if index[0] * coef + intercept <= index[1]:
                to_return[index] = 1
    return to_return


def get_band(drum, coef, intercept, bandwidth):
    """Calculate a ribbon path on the drum"""
    to_return = np.zeros(drum.shape)
    t1 = line_mask(drum, coef, intercept + bandwidth / 2, upper=True)
    t2 = line_mask(drum, coef, intercept - bandwidth / 2, upper=False)
    to_return = t1 + t2
    return np.where(to_return == 2, 1, 0)

single_band = get_band(drum, 1 / 10, 130, bandwidth=15)

# Visualize the result !
plt.imshow(single_band)
plt.show()

Numba 为我的代码创造了奇迹，将运行时间从 5.8 秒减少到 86 毫秒（特别感谢@Maarten-vd-Sande）：

from numba import jit
@jit(nopython=True, parallel=True)
def line_mask(drum, coef, intercept, upper=True, accuracy=accuracy):
    ...

仍然欢迎更好的 numpy 解决方案 ;-)

Answer 1

这里根本不需要任何循环。您实际上有两个不同的 line_mask 函数。两者都不需要显式循环，但你可能会通过在 if 和 else 中使用一对 for 循环而不是 if 重写它来获得显着的加速] 和 else 在 for 循环中，它被计算了很多次。

真正的 numpythonic 做法是正确矢量化您的代码以在没有任何循环的情况下对整个数组进行操作。这是 line_mask:

的矢量化版本

def line_mask(drum, coef, intercept, upper=True, accuracy=accuracy):
    """Masks a half of the array"""
    r = np.arange(drum.shape[0]).reshape(-1, 1)
    c = np.arange(drum.shape[1]).reshape(1, -1)
    comp = c.__lt__ if upper else c.__ge__
    return comp(r * coef + intercept)

设置r和c的形状为(m, 1)和(n, 1)，结果为(m, n)，称为broadcasting, 并且是 numpy 中矢量化的主要内容。

更新的 line_mask 的结果是一个布尔掩码（顾名思义）而不是一个浮点数组。这使它更小，并有望完全绕过浮动操作。您现在可以重写 get_band 以使用掩码而不是加法：

def get_band(drum, coef, intercept, bandwidth):
    """Calculate a ribbon path on the drum"""
    t1 = line_mask(drum, coef, intercept + bandwidth / 2, upper=True)
    t2 = line_mask(drum, coef, intercept - bandwidth / 2, upper=False)
    return t1 & t2

程序的其余部分应该保持不变，因为这些函数保留了所有接口。

如果需要，您可以用三行（仍然有些清晰）重写大部分程序：

coeff = 1/10
intercept = 130
bandwidth = 15

r, c = np.ogrid[:drum.shape[0], :drum.shape[1]]
check = r * coeff + intercept
single_band = ((check + bandwidth / 2 > c) & (check - bandwidth /  2 <= c))

向量化 numpy 中的 for 循环以计算胶带重叠

Vectorize a for loop in numpy to calculate duct-tape overlaping

python

for-loop

numpy

vectorization

numba