python - 在 numpy 中矢量化 for 循环以计算胶带重叠

Question

我正在使用 python 创建一个应用程序来计算胶带重叠(对分配器进行建模，将产品应用到旋转的滚筒上)。

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

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职能。两者都不需要显式循环，但您可能会通过使用一对 for 重写它来获得显着的加速。 if 中的循环和 else ，而不是 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是一个 bool 掩码(顾名思义)而不是一个浮点数组。这使它更小，并有望完全绕过 float 操作。您现在可以重写 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

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

如果你愿意，你可以用三行(仍然有点清晰)重写你的大部分程序:

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))