Search code examples
pythonnumpyscipyvectorizationcosine-similarity

How to vectorize an operation between a 2D array and another vector?


I have a numpy 2D array like this:

[[1, 2], [3, 4]]

And a separate vector like this: [5, 6].

The operation in this case is np.inner, as part of my larger task of taking the cosine similarity between each row of the 2D array and the separate vector.

My expected output is [np.inner([1, 2], [5, 6]), np.inner([3, 4], [5, 6]]. I can accomplish this using apply_along_axis, but is there any way to vectorize this operation and make it more performant than apply_along_axis?

There are lots of answered questions regarding how to vectorize these types of operations between pairs of rows of two 2D arrays, but in this situation I need to vectorize the operation between one 2D array and another vector. I could turn [5, 6] into [[5, 6], [5, 6]] and then have the vectorization that way, but at large scales I need a solution where I can use the other vector itself in the operation rather than turning it into a 2D array with a bunch of row copies.


Solution

  • Sound like you need `np.dot'

    In [3]: a = np.array([[1, 2], [3, 4]])
    
    In [4]: b = np.array([5, 6])
    
    In [5]: a.dot(b)
    Out[5]: array([17, 39])
    
    In [6]: np.inner(a[0], b)
    Out[6]: 17
    
    In [7]: np.inner(a[1], b)
    Out[7]: 39