Python 如果满足另一个数组中的值条件，则求numpy数组中的值之和

我面临一个问题，就是如何对函数进行矢量化，以便它能够有效地应用于numpy数组

我的节目条目：

• Nb_粒子线的二维阵列，3列（基本x、y、z坐标，只有z与困扰我的部分相关），Nb_粒子可以多达数十万个
• 带有Nb_粒子值的prop_部分1D数组。这一部分我得到了涵盖，创造了一些不错的numpy函数；我只是在这里放了一个基本的分布，它反映了真实的值
• 一个z_距离1D数组，一个简单的np.arange，介于z=0和z=z_max之间

然后是需要时间的计算，因为我找不到一种只对数组执行numpy操作的方法。我想做的是：

• 对于z_距离中的所有距离z_i，如果相应的粒子坐标z_粒子属性部分中的所有值之和。这将返回与z_距离相同长度的1D数组

到目前为止，我的想法是：

• 版本0，for loop、enumerate和np.where检索需要求和的值的索引。显然相当长
• 版本1，在新阵列上使用遮罩（z坐标和粒子属性的组合），并在遮罩阵列上求和。似乎比v0好
• 版本2，另一个掩码和一个np.vectorize，但我知道它并不有效，因为vectorize基本上是for循环。仍然比v0好
• 版本3，我尝试在一个可以直接应用于z_距离的函数上使用mask，但到目前为止它还不起作用

所以，我来了。这可能与排序和累计总数有关，但我不知道如何做到这一点，因此任何帮助都将不胜感激。请查找下面的代码，以使事情更清楚

提前谢谢

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)
# =============================

将numpy导入为np
导入时间
将matplotlib.pyplot作为plt导入
#粒子位置的创建
Nb_零件=150_000
位置部分=10*np.random.rand（Nb部分，3）
位置部件[：，0]=位置部件[：，1]=0
#有用的属性创建
β=1/1.5
项目部分=（1/beta）*np.exp（-pos\u部分[：，2]/beta）
z_距离=np.arange（0,10,0.1）
#我的版本0
t0=时间。时间（）
结果=np.空（len（z_距离））
对于索引距离，枚举中的值距离（z距离）：
positions=np.where（pos_part[：，2]完全用numpy编写第一个版本可以获得更高的性能。用
np.sum
替换pythons
sum
。而不是
对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

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)