问题描述
我有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,0] #number of start node for each pipe,0<=i<Nj
j=[1,5] #number of end node for each pipe,0<=j<Nj
b=[1.3,0.5,1.5,2.7,2.2,3.1] #value for each pipe,Np elements
node=[0,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,-3.6]
在我自己的管网中,我有Nj = 150628和Np = 157040,因此我创建的循环花费了太多时间(大约0.6s)。所以我想问我如何向量化它?谢谢! 我尝试执行以下矢量化代码:
F=np.zeros(Nj)
F[i]=F[i]+b
F[j]=F[j]-b
F=[ 3.1,0.,-3.1,0. ]
这给出了错误的结果,因为在给定节点的开始节点或结束节点处可能有多个管道,但是在任一侧仅计数其中一个。 另外,如果我创建两个稀疏矩阵Mat_i和Mat_j来代表连接到起始节点/结束节点的所有管道,然后对其进行迭代,会更快吗? (我正在使用python 3.7)。
我刚刚设法使它工作:
F=np.bincount(i,weights=b,minlength=Nj)-np.bincount(j,minlength=Nj)
我也愿意使用Numba,就像我在代码的另一部分中使用@njit(parallel=True)
来实现标量函数一样,它通过使用我的cpu的所有8个线程大大加快了工作速度。
解决方法
使用纯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))