python - 嵌套循环中的向量化操作 : 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。