[英]Optimize operations with arrays in numpy
我必須應用一些我用python編寫的數學公式:
for s in range(tdim):
sum1 = 0.0
for i in range(dim):
for j in range(dim):
sum1+=0.5*np.cos(theta[s]*(i-j))*
eig1[i]*eig1[j]+eig2[i]+eig2[j])-0.5*np.sin(theta[s]*(i-j))*eig1[j]*eig2[i]-eig1[i]*eig2[j])
PHi2.append(sum1)
現在,這是正確的,但顯然效率很低,另一種方法是:
for i in range(dim):
for j in range(dim):
PHi2 = 0.5*np.cos(theta*(i-j))*(eig1[i]*eig1[j]+eig2[i]+eig2[j])-0.5*np.sin(theta*(i-j))*(eig1[j]*eig2[i]-eig1[i]*eig2[j])
但是,第二個示例在PHi2的所有元素中給了我相同的數字,因此這速度更快,但答案是錯誤的。 您如何正確,更有效地執行此操作?
注意:eig1和eig2具有相同的尺寸d,θ和PHi2具有相同的尺寸D,但d!= D。
您可以使用蠻力廣播方法,但是您正在創建形狀為(D, d, d)
的中間陣列,如果陣列的大小適中,則該陣列可能會失控。 此外,在使用沒有改進的廣播時,您需要從最內層循環重新計算很多計算,而您只需要執行一次即可。 如果您首先為i - j
所有可能值計算必要的參數並將它們相加,則可以在外循環上重用這些值,例如:
def fast_ops(eig1, eig2, theta):
d = len(eig1)
d_arr = np.arange(d)
i_j = d_arr[:, None] - d_arr[None, :]
reidx = i_j + d - 1
mult1 = eig1[:, None] * eig1[ None, :] + eig2[:, None] + eig2[None, :]
mult2 = eig1[None, :] * eig2[:, None] - eig1[:, None] * eig2[None, :]
mult1_reidx = np.bincount(reidx.ravel(), weights=mult1.ravel())
mult2_reidx = np.bincount(reidx.ravel(), weights=mult2.ravel())
angles = theta[:, None] * np.arange(1 - d, d)
return 0.5 * (np.einsum('ij,j->i', np.cos(angles), mult1_reidx) -
np.einsum('ij,j->i', np.sin(angles), mult2_reidx))
如果我們將M4rtini的代碼重寫為比較函數:
def fast_ops1(eig1, eig2, theta):
d = len(eig1)
D = len(theta)
s = np.array(range(D))[:, None, None]
i = np.array(range(d))[:, None]
j = np.array(range(d))
ret = 0.5 * (np.cos(theta[s]*(i-j))*(eig1[i]*eig1[j]+eig2[i]+eig2[j]) -
np.sin(theta[s]*(i-j))*(eig1[j]*eig2[i]-eig1[i]*eig2[j]))
return ret.sum(axis=(-1, -2))
我們組成一些數據:
d, D = 100, 200
eig1 = np.random.rand(d)
eig2 = np.random.rand(d)
theta = np.random.rand(D)
速度提高非常明顯,在原始代碼的115倍的基礎上提高了80倍,從而使速度提高了9000倍:
In [22]: np.allclose(fast_ops1(eig1, eig2, theta), fast_ops(eig1, eig2, theta))
Out[22]: True
In [23]: %timeit fast_ops1(eig1, eig2, theta)
10 loops, best of 3: 145 ms per loop
In [24]: %timeit fast_ops(eig1, eig2, theta)
1000 loops, best of 3: 1.85 ms per loop
這通過廣播起作用。
對於tdim = 200 and dim = 100
。
14秒與原始。
版本為120毫秒。
s = np.array(range(tdim))[:, None, None]
i = np.array(range(dim))[:, None]
j = np.array(range(dim))
PHi2 =(0.5*np.cos(theta[s]*(i-j))*(eig1[i]*eig1[j]+eig2[i]+eig2[j])-0.5*np.sin(theta[s]*(i-j))*(eig1[j]*eig2[i]-eig1[i]*eig2[j])).sum(axis=2).sum(axis=1)
在代碼的第一位,您有0.5*np.cos(theta[s]*(ij))...
但在第二位是0.5*np.cos(theta*(ij))...
除非您為第二部分代碼定義了不同的theta,否則很可能是造成麻煩的原因。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.