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

Question

我在矢量化函数以使其有效地应用于 numpy 数组时遇到问题。

我的节目条目：

• Nb_particles 行的 pos_part 2D 数组，3 列（基本上是 x、y、z 坐标，只有 z 与困扰我的部分相关） Nb_particles 可以达到数十万。
• 具有 Nb_particles 值的 prop_part 一维数组。这部分我已经介绍过了，创建是用一些不错的 numpy 函数进行的；我只是在这里放了一个类似于真实值的基本分布。
• z_distances 一维数组，一个简单的 np.arange，介于 z=0 和 z=z_max 之间。

然后是需要时间的计算，因为我无法找到一种仅通过数组的 numpy 操作来正确处理事情的方法。我想做的是：

• 对于 z_distances 中的所有距离 z_i，如果对应的粒子坐标 z_particle prop_part 中的所有值相加。这将返回一个与 z_distances 长度相同的一维数组。

到目前为止我的想法：

• 版本 0、for 循环、enumerate 和 np.where 确实检索了我需要求和的值的索引。显然很长。
• 版本 1，在新阵列上使用蒙版（z 坐标和粒子属性的组合），并在蒙版阵列上求和。似乎比 v0 好
• 版本 2，另一个掩码和一个 np.vectorize，但我知道它效率不高，因为矢量化基本上是一个 for 循环。似乎仍然比 v0 好
• 第 3 版，我正在尝试在可以直接应用于 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 编写您的第一个版本，您可以获得更多性能。将 pythons 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 通常会产生一些开销，我们可以通过为一小部分`z_distances 运行函数来消除这些开销。 下面的代码在我的笔记本电脑上实现了大约两倍于纯 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)