时间戳

首页 > 解决方案 > 如果满足另一个数组中值的条件，则对 numpy 数组中的值求和

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

问题描述

我面临一个矢量化函数的问题，以便它有效地应用于 numpy 数组。

我的节目条目：

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

到目前为止我的想法：

所以，我在这里。可能与排序和累积总和有关，但我不知道该怎么做，所以任何帮助将不胜感激。请在下面找到代码以使事情更清楚

提前致谢。

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

标签： pythonnumpynumpy-ndarray

解决方案

通过完全用 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通常会产生一些开销，我们可以通过为一小部分`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)

推荐阅读