Numpy 2D dot operation, indexing and iteration. How to optimize?

Question

I need to remove this bottleneck from the code. Hope someone can help me.

I have the following iteration (extract from a nn inverse distance interpolation algorithm):

jinterpol=0
for w, ix in zip(wds, ixs):
        wz = np.vdot(w, z[ix])
        result[jinterpol] = wz
        jinterpol += 1

I've tried with removing the for loop:

result = np.dot(wds, z[ixs])

but of course it says ValueError: objects are not aligned

Can you give me some tips? Many thanks.

Shapes are:

• wds: (550800, 8)
• z: (212065,)
• z[ixs]: (550800, 8)
• ixs:(550800, 8)

and

• len(zip(distances, ixs)): 550800
• w: (8,)
• ix: (8,)

Answer 1

You can use np.einsum :

>>> import numpy as np
>>> wds = np.random.rand(550800, 8)
>>> z = np.random.rand(212065)
>>> ixs = np.random.randint(212065, size=(550800, 8))

>>> np.einsum('ij,ij->i', wds, z[ixs])
array([ 1.65069924,  3.26203701,  3.16035664, ...,  1.76963986,
        2.09727537,  1.94905991])

>>> np.vdot(wds[0], z[ixs[0]])
1.6506992361953157
>>> np.vdot(wds[1], z[ixs[1]])
3.2620370116548827