嵌套循环中的向量化操作：Python

Question

我有 2D numpy 数组，我需要两个嵌套循环来迭代它的每个元素。我想对嵌套循环进行一些矢量化，但我一直收到错误提示，

j= np.arange (0,x.shape [1])
IndexError: tuple index out of range

这些是原始的嵌套循环：

for k in range(A.shape[0]):
        for j in range(A.shape[1]):
            A[k,j] = method1(x[k],x[j],a,c0,c1)

这就是我尝试根据这个答案进行矢量化的方法，https://codereview.stackexchange.com/questions/17702/python-numpy-running-15x-slower-than-matlab-am-i-using-numpy-effeciently，

j= np.arange (0, A.shape [1])
    for k in range(A.shape[0]):

            A[k,j] = method1(x[k],x[j],a,c0,c1)

我尝试更改 np.arange 中的索引，但没有成功。

谁能告诉我如何解决这个问题？

谢谢。

编辑：通过@ajcr 评论，我的错误在于这一行，j= np.arange (0,x.shape [1])，我应该使用二维数组“A”的列数作为 j = np.arange(0,A.shape [1])，但我错误地使用了一维数组x，因此出现错误。它现在工作得很好。

Answer 1

由于您不提供函数本身，因此很难给您答案。如果可以将您的计算矢量化，这将是可行的方法。否则，您可以使用np.vectorize：

import numpy as np

def my_fun(x,y,a,b):
    if x > y:
        return a*x**2 + y
    else:
        return b*x**2 + a*y

vec_fun = np.vectorize(my_fun)
x = np.random.rand(100000)
y = np.random.rand(100000)

%%timeit
for xx,yy in zip(x,y):
    my_fun(xx,yy,1,2)

10 个循环，3 个循环中的最佳值：每个循环 138 毫秒

%%timeit
    vec_fun(x,y,1,2)

10 个循环，3 个循环中的最佳值：每个循环 65.4 毫秒

如果您只是循环输入参数列表，您可以使用multiprocessing.Pool 来利用您的所有内核。

编辑：正如 BlackCat 所指出的， np.vectorize 函数可能不会加快速度。如果是这种情况，我能想到的唯一解决方案是使用所有内核或切换到 Cython。