CODEWITHSUNDEEP
pythonarraysnumpymatrixstatistics
Viktor VN
Viktor VN

Reputation: 235

Calculate a dot product in bulk with numpy

I am trying to calculate the covariance matrices of a coordinate transformation (from spherical to Cartesian coordinates). The dataset contains 504 different spherical coordinates.

I calculate the covariance matrix of the transformation using an error propagation which i implemented using a matrix product.

This matrix product is A dot CR dot B. The general shape of these matrices can be seen here:

A = [[dx/dr, dx/dt, dx/dp]
     [dy/dr, dy/dt, dy/dp]
     [dz/dr, dz/dt, dz/dp]]

B = [[dx/dr, dy/dr, dz/dr]
     [dx/dt, dy/dt, dz/dt]
     [dx/dp, dy/dp, dz/dp]]

CR = [[0.001 ** 2, 0, 0]
      [0, 0.003 ** 2, 0]
      [0, 0, 0.005 ** 2]]

Now i want to implement this in a way which doesn't use for loops for performance reasons. A and B in the actual code (at the bottom of the post) are arrays containing 504, 3x3 matrices. The shape of the array is thus (508, 3, 3). What I want to do is take the dot product described above for all corresponding matrices A[0] dot CR dot B[0], A[0] dot CR dot B[0], ... and have the resulting array be of shape (508, 3, 3). I currently have implemented this using a list comprehension (see calculation of CX), but I want to remove the for loops entirely.

I tried passing the (508, 3, 3) arrays to np.dot() -> CX = np.dot(np.dot(A, CR), B) but this results in a (508, 3, 508, 3) array which is not what i want. There doesn't seem to be a keyword argument to specifying which axes it should use for the dot product like there is for np.transpose(). Does anyone know how I can implement this?

SIGMA_R, SIGMA_THETA, SIGMA_PHI = 0.001, 0.003, 0.005
CR = np.array([[SIGMA_R ** 2, 0, 0],
               [0, SIGMA_THETA ** 2, 0],
               [0, 0, SIGMA_PHI ** 2]])

A = np.transpose(np.array(([derivative_r(r, theta, phi),
                            derivative_theta(r, theta, phi),
                            derivative_phi(r, theta, phi)])))

B = np.transpose(A, axes=(0, 2, 1))
CX = np.array([np.dot(np.dot(A[i], CR), B[i]) for i in range(504)])

Upvotes: 0

Views: 119

Answers (1)

Daniel F
Daniel F

Reputation: 14399

Looks like you need np.einsum

CX = np.einsum('ijk, kl, ilm -> ijm', A, CR, B)

Or even:

CX = np.einsum('ijk, kl, iml -> ijm', A, CR, A)

Upvotes: 1

Related Questions