如果滿足另一個數組中值的條件，則對 numpy 數組中的值求和

Question

我在對 function 進行矢量化以使其有效地應用於 numpy 陣列時遇到問題。

我的節目條目：

• Nb_particles線的pos_part 2D 陣列，3 列（基本上是 x,y,z 坐標，只有 z 與困擾我的部分相關） Nb_particles可以多達幾十萬。
• 具有Nb_particles值的prop_part一維數組。 這部分我已經介紹過了，創建是用一些不錯的 numpy 函數； 我只是在這里放了一個類似於真實值的基本分布。
• 一個z_distances 1D 數組，一個簡單的 np.arange 介於 z=0 和 z=z_max 之間。

然后是需要時間的計算，因為我無法找到一種方法來正確處理 arrays 的 numpy 操作。 我想做的是：

• 對於 z_distances 中的所有距離z_i ，如果對應的粒子坐標 z_particle < z_i ，則將prop_part 中的所有值相加。 這將返回一個與z_distances長度相同的一維數組。

到目前為止我的想法：

• 版本 0、 for循環、枚舉和np.where確實檢索我需要求和的值的索引。 顯然很長。
• 版本 1，在新數組上使用遮罩（z 坐標和粒子屬性的組合），並在遮罩數組上求和。 好像比v0好
• 版本 2，另一個掩碼和一個np.vectorize ，但我知道它效率不高，因為 vectorize 基本上是一個for循環。 似乎仍然比 v0 好
• 版本 3，我正在嘗試在 function 上使用掩碼，我可以直接將其應用於 z_distances，但到目前為止還沒有工作。

所以，我在這里。 可能與排序和累積總和有關，但我不知道該怎么做，所以任何幫助將不勝感激。 請在下面找到代碼以使事情更清楚

提前致謝。

import numpy as np
import time
import matplotlib.pyplot as plt

# Creation of particles' positions
Nb_part = 150_000
pos_part = 10*np.random.rand(Nb_part,3)
pos_part[:,0] = pos_part[:,1] = 0

#usefull property creation
beta = 1/1.5
prop_part = (1/beta)*np.exp(-pos_part[:,2]/beta)
z_distances = np.arange(0,10,0.1)


#my version 0
t0=time.time()
result = np.empty(len(z_distances))
for index_dist, val_dist in enumerate(z_distances):
    positions = np.where(pos_part[:,2]<val_dist)[0]
    result[index_dist] = sum(prop_part[i] for i in positions)
print("v0 :",time.time()-t0)

#A graph to help understand
plt.figure()
plt.plot(z_distances,result, c="red")
plt.ylabel("Sum of particles' usefull property for particles with z-pos<d")
plt.xlabel("d")


#version 1 ??
t1=time.time()
combi = np.column_stack((pos_part[:,2],prop_part))
result2 = np.empty(len(z_distances))
for index_dist, val_dist in enumerate(z_distances):
    mask = (combi[:,0]<val_dist)
    result2[index_dist]=sum(combi[:,1][mask])
print("v1 :",time.time()-t1)
plt.plot(z_distances,result2, c="blue")


#version 2
t2=time.time()
def themask(a):
    mask = (combi[:,0]<a)
    return sum(combi[:,1][mask])
thefunc = np.vectorize(themask)
result3 = thefunc(z_distances)
print("v2 :",time.time()-t2)
plt.plot(z_distances,result3, c="green")


### This does not work so far
# version 3
# =============================
# t3=time.time()
# def thesum(a):
#     mask = combi[combi[:,0]<a]
#     return sum(mask[:,1])
# result4 = thesum(z_distances)
# print("v3 :",time.time()-t3)
# =============================

Answer 1

通過完全在 numpy 中編寫您的第一個版本，您可以獲得更多的性能。 將 python 的sum替換為np.sum 。 而不是for i in positions列表理解中，只需傳遞您正在創建的positions掩碼。 事實上， np.where不是必需的，我最好的版本如下所示：

#my version 0
t0=time.time()
result = np.empty(len(z_distances))
for index_dist, val_dist in enumerate(z_distances):
    positions = pos_part[:, 2] < val_dist
    result[index_dist] = np.sum(prop_part[positions])
print("v0 :",time.time()-t0)
# out: v0 : 0.06322097778320312

如果 z_distances 很長，使用numba可以更快一些。 第一次運行calc通常會產生一些開銷，我們可以通過運行 function 來消除一些開銷。 下面的代碼在我的筆記本電腦上實現了大約兩倍於純 numpy 的加速。

import numba as nb
@nb.njit(parallel=True)
def calc(result, z_distances):
    n = z_distances.shape[0]
    for ii in nb.prange(n):
        pos = pos_part[:, 2] < z_distances[ii]
        result[ii] = np.sum(prop_part[pos])
    return result

result4 = np.zeros_like(result)
# _t = time.time()
# calc(result4, z_distances[:10])
# print(time.time()-_t)
t3 = time.time()
result4 = calc(result4, z_distances)
print("v3 :", time.time()-t3)
plt.plot(z_distances, result4)