我想在第二個軸上執行一個二維陣列的總和，但在一個可變的范圍內。未矢量化它是：`
import numpy as np
nx = 3
ny = 5
a = np.ones((nx, ny))
left_bnd = np.array([0, 1, 0])
right_bnd = np.array([2, 2, 4])
b = np.zeros(nx)
for jx in range(nx):
b[jx] = np.sum(a[jx, left_bnd[jx]: right_bnd[jx]])
print(b)
輸出 b 為 [2。1. 4.] 我很想對回圈進行矢量化，有點
b = np.sum(a[:, left_bnd[:]: right_bnd[:], axis=1)
加快計算速度，因為我的“n”通常是幾個 1e6。不幸的是，我找不到合適的作業語法。
uj5u.com熱心網友回復：
在 for 回圈中手動求和的 jittednumba實作大約快 100 倍。在函式內部使用np.sumwith slicing 的numba速度只有一半。此解決方案假定所有切片都在有效范圍內。
為基準測驗生成足夠大的樣本資料
import numpy as np
import numba as nb
np.random.seed(42) # just for reproducibility
n, m = 5000, 100
a = np.random.rand(n,m)
bnd_l, bnd_r = np.sort(np.random.randint(m 1, size=(n,2))).T
與numba. 請確保通過運行該函式至少兩次來對已編譯的熱代碼進行基準測驗。
@nb.njit
def slice_sum(a, bnd_l, bnd_r):
b = np.zeros(a.shape[0])
for j in range(a.shape[0]):
for i in range(bnd_l[j], bnd_r[j]):
b[j] = a[j,i]
return b
slice_sum(a, bnd_l, bnd_r)
輸出
# %timeit 1000 loops, best of 5: 297 μs per loop
array([ 4.31060848, 35.90684722, 38.03820523, ..., 37.9578962 ,
3.61011028, 6.53631388])
在numpypython 回圈中（這是一個很好、簡單的實作）
b = np.zeros(n)
for j in range(n):
b[j] = np.sum(a[ j, bnd_l[j] : bnd_r[j] ])
b
輸出
# %timeit 10 loops, best of 5: 29.2 ms per loop
array([ 4.31060848, 35.90684722, 38.03820523, ..., 37.9578962 ,
3.61011028, 6.53631388])
驗證結果是否相等
np.testing.assert_allclose(slice_sum(a, bnd_l, bnd_r), b)
uj5u.com熱心網友回復：
這是一個純 numpy 解決方案，它接近發布的 numba 解決方案的速度。它利用reduceat了，但設定非常復雜。
def slice_sum_np(a, left_bnd, right_bnd):
nx, ny = a.shape
linear_indices = np.c_[left_bnd, right_bnd] ny * np.arange(nx)[:,None]
sums = np.add.reduceat(a.ravel(), linear_indices.ravel())[::2]
# account for reduceat special case
sums[left_bnd >= right_bnd] = 0
return sums
轉載請註明出處，本文鏈接：https://www.uj5u.com/yidong/453074.html
上一篇：如何添加和使用“Numpy”庫？
下一篇：Numpy，熊貓合并行