CODEWITHSUNDEEP
pythonarraysnumpy
k0ngcrete
k0ngcrete

Reputation: 116

Get the components of a multidimensional array dot product without a loop

I want to vectorise the dot product of several 3x3 matrices (rotation matrix around x-axis) with several 3x1 vectors. The application is the transformation of points (approx 500k per array) from one to another coordinate system.
Here in the example only four of each. Hence, the result should be again 4 times a 3x1 vector, respectively the single components x,y,z be a 4x0 vector. But I cannot get the dimensions figured out: Here the dot product with tensordot in results in a shape of (4,3,4), of which I need the diagonals again:

x,y,z = np.zeros((3,4,1))
rota = np.arange(4* 3 * 3).reshape((4,3, 3))
v= np.arange(4 * 3).reshape((4, 3))
result = np.zeros_like(v, dtype = np.float64)

vec_rotated = np.tensordot(rota,v, axes=([-1],[1]))

for i in range(result.shape[0]):
    result[i,:] = vec_rotated[i,:,i]
x,y,z = result.T

How can i vectorise the complete thing?

Upvotes: 0

Views: 47

Answers (2)

Divakar
Divakar

Reputation: 221614

Use np.einsum for an efficient solution -

x,y,z = np.einsum('ijk,ik->ji',rota,v)

Alternative with np.matmul/@ operator in Python 3.x -

x,y,z = np.matmul(rota,v[:,:,None])[...,0].T
x,y,z = (rota@v[...,None])[...,0].T

Upvotes: 1

k0ngcrete
k0ngcrete

Reputation: 116

works via transpose to obtain one component per diagonal:

vec_rotated = vec_rotated.transpose((1,0,2))
x,y,z = np.diag(vec_rotated[0,:,:]),np.diag(vec_rotated[1,:,:]),np.diag(vec_rotated[2,:,:])

Upvotes: 0

Related Questions