我正在使用Python创建一个应用程序,用于计算胶带重叠(模拟旋转鼓上的产品应用)。
我有一个可以正常工作的程序,但速度非常慢。我正在寻求优化“for”循环以填充numpy数组的解决方案。有人能帮我矢量化下面的代码吗?
Numba 对我的代码产生了神奇的作用,将运行时间从 5.8 秒降低到 86 毫秒(特别感谢 @Maarten-vd-Sande)。
一个更好的解决方案仍然欢迎使用numpy;-)
我有一个可以正常工作的程序,但速度非常慢。我正在寻求优化“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;-)