Python - 如何加速 for 循環從另一個 numpy 數組計算創建一個 numpy 數組

Question

首先，為模糊的標題道歉，我想不出這個問題的合適名稱。

我有以下格式的 3 個 numpy 數組：

N = ([[13, 14, 15], [2, 5, 7], [4, 6, 8] ... 幾十萬個元素長

e1 = [1, 0, 0]

e2 = [0, 1, 0]

這個想法是創建第四個數組“v”，它應該具有與“N”相同的維度，但將根據 if 語句給出值。 這是我目前擁有的應該更好地解釋問題的內容：

v = np.zeros([len(N), 3])    

for i in range(0, len(N)):
    if((N*e1)[i,0] != 0):
        v[i] = np.cross(N[i],e1)
    else:
        v[i] = np.cross(N[i],e2)

此代碼執行我的要求，但執行的時間比預期的要長（> 5 分鍾）。 是否有任何形式的列表理解或類似的概念可以用來提高代碼的效率？

Answer 1

您可以使用numpy.where替換 if-else 並使用廣播矢量化進程，這里是numpy.where一個選項：

import numpy as np
np.where(np.repeat(N[:,0] != 0, 3).reshape(1000,3), np.cross(N, e1), np.cross(N, e2))

這里的一些基准：

1)數據設置：

N = np.array([np.random.randint(0,10,3) for i in range(1000)])
N

#array([[3, 5, 0],
#       [5, 0, 8],
#       [4, 6, 0],
#       ..., 
#       [9, 4, 2],
#       [6, 9, 3],
#       [2, 9, 2]])

e1 = np.array([1, 0, 0])
e2 = np.array([0, 1, 0])

2)時間：

def forloop():
    v = np.zeros([len(N), 3]);    
​
    for i in range(0, len(N)):
        if((N*e1)[i,0] != 0):
            v[i] = np.cross(N[i],e1)
        else:
            v[i] = np.cross(N[i],e2)
    return v

def forloop2():
    v = np.zeros([len(N), 3])    
​
    # Only calculate this one time.
    my_product = N*e1
​
    for i in range(0, len(N)):
        if my_product[i,0] != 0:
            v[i] = np.cross(N[i],e1)
        else:
            v[i] = np.cross(N[i],e2)               
    return v

%timeit forloop()
10 loops, best of 3: 25.5 ms per loop

%timeit forloop2()
100 loops, best of 3: 12.7 ms per loop    

%timeit np.where(np.repeat(N[:,0] != 0, 3).reshape(1000,3), np.cross(N, e1), np.cross(N, e2))
10000 loops, best of 3: 71.9 µs per loop

3)所有方法的結果檢查：

v1 = forloop()   

v2 = np.where(np.repeat(N[:,0] != 0, 3).reshape(1000,3), np.cross(N, e1), np.cross(N, e2))

v3 = forloop2()

(v3 == v1).all()
# True

(v1 == v2).all()
# True

Answer 2

我不確定你想要做什么，但我知道為什么這個特定的代碼對你來說這么慢。 最嚴重的違規者是(N*e1) 。 這是一個簡單的計算，它使用 numpy 運行得非常快，但是您在循環內執行它len(N)次！

通過將代碼拉到循環之外，我能夠在不到 15 秒的時間N == 1000000我的機器上使用N == 1000000執行您的代碼。 下面舉例。

v = np.zeros([len(N), 3])    

# Only calculate this one time.
my_product = N*e1

for i in range(0, len(N)):
    if my_product[i,0] != 0):
        v[i] = np.cross(N[i],e1)
    else:
        v[i] = np.cross(N[i],e2)

另一個答案演示了如何避免 for 循環和 if 語句以提高速度，但代價是代碼可讀性稍差。