dimanche 6 décembre 2015

Speed up numpy for and if-else by vectorizing?

Recently, I discovered the bottleneck (slow-down - especially with a large 3D array input) came from a section of the code as such:

for i in range(0, h):
    for j in range(0, w):
        for k in range(0, c):
            if k == 0:
                u = 2
            elif k==2:
                u = 0
            else:
                u = 1

            if ra[i,j,k] < cr[0, u]:
                im[i, j, k] = 0
            elif ra[i,j,k] > cr[200, u]:
                im[i, j, k] = 200
            else:
                p = np.argmin(np.abs(ra[i,j,k]-cr[:,u]))

                if ra[i,j,k] >= cr[p,u]:
                    p1 = p
                    p2 = p+1
                elif ra[i,j,k] < cr[p,u]:
                    p2 = p
                    p1 = p-1

                im[i,j,k] = (ra[i,j,k]-cr[p1,u])/(cr[p2,u]-cr[p1,u]) + p1 - 1

I have been thinking that the slow-down could be a result of having to access individual elements, performing if-else conditional check in each step.

Can I check if I could vectorise this whole fragment of code, instead of the element-wise check method now? (meaning to compress the code as much as possible in memory, so I can get better speed improvement while still doing the same operations)

I have thought of np.where(ra[i,j,k] < cr[0,u], 0, 255) or the use of a ternary operator im[i,j,k] = 0 if rad[i,j,k] < crf[0,u] else 255.... but I am at a lost as to extending it to if-elif-else.

Aucun commentaire:

Enregistrer un commentaire