时间戳

首页 > 解决方案 > 在numpy或python中将一个向量广播到另一个不同大小的向量中

python - 在numpy或python中将一个向量广播到另一个不同大小的向量中

问题描述

我有Np链接（管道）和Nj连接点（节点）。每个链接k都有一个起始节点i和结束节点j，链接值为b。b我想通过添加给定节点是否是链接的起始节点并减去b给定节点是否是链接的结尾来计算链接对每个节点的贡献。我的代码是这样的：

for k in range(Np):   #iterate over pipes
    F[i[k]]=F[i[k]]+b[k]   #add value of b if pipe k is at the start of the node
    F[j[k]]=F[j[k]]-b[k]   #substract b if pipe k is at the end of the node

用于运行此代码的数据示例：

Np=7   #number of pipes
Nj=6  #number of junctions (nodes)
pipe=[0,1,2,3,4,5,6]  #number of pipes are consecutive, k=0:Np
i=[0,1,2,3,3,3,0]     #number of start node for each pipe, 0<=i<Nj
j=[1,5,1,1,2,4,5]     #number of end node for each pipe, 0<=j<Nj
b=[1.3,0.5,1.5,2.7,1.5,2.2,3.1]   #value for each pipe, Np elements
node=[0,1,2,3,4,5]   #node numbers are consecutive, 0:Nj
F=np.zeros(Nj)   #intializing F, size is equal to number of nodes

循环的结果是：

F=[+1.3+3.1,+0.5-1.3-1.5-2.7,+1.5-1.5,+2.7+1.5+2.2,-2.2,-0.5-3.1]

或者

F=[4.4, -5.0, 0.0, 6.4, -2.2, -3.6]

在我自己的管道网络中，我有 Nj=150628 和 Np=157040，所以我创建的循环花费了太多时间（大约 0.6 秒）。所以我想问我如何矢量化它？谢谢！我尝试执行以下矢量化代码：

F=np.zeros(Nj)
F[i]=F[i]+b
F[j]=F[j]-b
F=[ 3.1, -2.2,  0. ,  2.2, -2.2, -3.1,  0. ]

这给出了错误的结果，因为可能有多个管道位于给定节点的开始或结束节点，但它只计算任一侧的管道。另外，如果我创建两个稀疏矩阵 Mat_i 和 Mat_j 代表连接到开始节点/结束节点的所有管道，然后迭代它会更快吗？（我使用的是 python 3.7）。

我只是设法让它与这个一起工作：

F=np.bincount(i,weights=b,minlength=Nj)-np.bincount(j,weights=b,minlength=Nj)

我也愿意使用 Numba，因为我@njit(parallel=True)在代码的另一部分中使用了标量函数，它通过使用我 CPU 的所有 8 个线程来加快速度。

标签： pythonnumpyscipynumbaarray-broadcasting

解决方案

使用纯 NumPy 执行此操作很棘手，因为您多次“选择”相同的索引（即，i和j中的值不是唯一的）。但是 Numba 将加速代码，基本上没有任何变化（我冒昧地使用+=简洁）：

@numba.njit
def smoke(F, b, i, j):
    for k in range(len(i)):   #iterate over pipes
        F[i[k]] += b[k]   #add value of b if pipe k is at the start of the node
        F[j[k]] -= b[k]   #subtract b if pipe k is at the end of the node

如果您的输入是列表，请像这样运行它：

smoke(F, np.asarray(b), np.asarray(i), np.asarray(j))

推荐阅读