CODEWITHSUNDEEP
pythonnumpyvectorization
FaceInvader
FaceInvader

Reputation: 89

NumPy - Dot Product along 3rd dimension without copying

I am trying to vectorize a function that takes as its input a 3-Component vector "x" and a 3x3 "matrix" and produces the scalar

def myfunc(x, matrix):
    return np.dot(x, np.dot(matrix, x))

However this needs to be called "n" times, and the vector x has different components each time. I would like to modify this function such that it takes as input some 3xn arrays (the columns of which are the vectors x) and produces a vector whose components are the scalars that would have been computed at each iteration.

I can write down an Einstein summation that does the job but it requires that I construct a 3x3xn stack of "copies" of the original 3x3. I am concerned that doing this will blow away any performance gains I get from trying to do this. Is there any way to compute the vector I want without making copies of the 3x3?

Upvotes: 1

Views: 258

Answers (1)

user2357112
user2357112

Reputation: 280778

Let x be the 3xN array and y be the 3x3 array. You're looking for

z = numpy.einsum('ji,jk,ki->i', x, y, x)

You also could have built that 3x3xN array you were talking about as a view of y to avoid copying, but it isn't necessary.

Upvotes: 3

Related Questions