FaceInvader - 1 year ago 104

Python Question

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?

Answer Source

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.