Broadcasting a vector into another vector of different size in numpy or python

user_14070764 0
Viewed 93

I have Np links (pipes) and Nj junctions (nodes). Each link k has a start node i and end node j, and a link value of b. I want to calculate the contribution of links to each node by adding b if the given node is the start node of a link, and substracting b if the given node is the end of a link. My code is something like this:

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

An example of the data used to run this code:

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

and the result of the loop would be:

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]

or

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

In my own pipe network I have Nj=150628 and Np=157040, so the loop I created takes way too much time (about 0.6s). So I wanted to ask how can I vectorize it? Thanks! I tried doing the following vectorized code:

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. ]

Which gives the wrong result because there might be multiple pipes that are at either start or end node of a given node, but it counts only one of them on either side. Also, would it be faster if I create two sparse matrices Mat_i and Mat_j that represent all pipes connected to start node / end node, and then iterate over that instead? (I'm using python 3.7).

I just managed to get it working with this:

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

I am also open to using Numba as I used @njit(parallel=True) in another part of my code for a scalar function & it sped things up alot by using all of the 8 threads of my CPU.

numba array-broadcasting scipy numpy python
1 Answers

Doing it with pure NumPy is tricky because you "select" the same index multiple times (i.e. the values in i and j are not unique). But Numba will accelerate the code with essentially no changes (I took the liberty of using += for concision):

@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

Run it like this if your inputs are lists:

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