CODEWITHSUNDEEP
pythonnumpy
pol
pol

Reputation: 53

Numpy vectorization for linear combination of numpy matrices

I have a numpy ndarray of shape (5,4,4) that is a set of 5 matrices 4x4. I would like to multiply that ndarray by a matrix of shape (3,5) and I would like to get a numpy ndarray of shape (3,4,4) where each matrix 4x4 in the result is the linear combination of the 5 4x4 matrices with the coefficients coming from the rows of the matrix (3,5).

I have the following, very basic, code:

import numpy as np


np.random.seed(10)
X = np.random.rand(5,4,4)
A = np.random.rand(3,5)

print( A*X )

which generates, a bit as I was expecting, the error 'operands could not be broadcast together with shapes (3,5) (5,4,4)'

I was reading a bit about broadcasting but I can't seem to find a way to write the operation in such a way it can be vectorized by Numpy.

Has anybody got a similar experience?

Upvotes: 2

Views: 115

Answers (2)

Lucifer
Lucifer

Reputation: 305

You can achieve the desired result by using the numpy dot product function np.dot() and reshaping the arrays appropriately. Here's an example code snippet that demonstrates how to perform the matrix multiplication and reshape the arrays:

import numpy as np

# Create the input arrays
array_5x4x4 = np.random.random((5, 4, 4))
array_3x5 = np.random.random((3, 5))

# Perform the matrix multiplication
result = np.dot(array_3x5, array_5x4x4.reshape(5, -1))

# Reshape the result array to the desired shape
result = result.reshape(3, 4, 4)

# Print the result
print(result)

In this example, array_5x4x4 is your original ndarray of shape (5, 4, 4), and array_3x5 is the matrix of shape (3, 5) containing the coefficients for the linear combination. The np.dot() function performs the matrix multiplication, and we use array_5x4x4.reshape(5, -1) to reshape array_5x4x4 to a 2D array with shape (5, 16), which makes the dot product compatible with array_3x5. Finally, we reshape the resulting array to the desired shape (3, 4, 4).

Upvotes: 1

Thomas Wagenaar
Thomas Wagenaar

Reputation: 6759

For these type of operations, np.einsum is perfect. Is the following what you want?

B = np.einsum('ij,jkl->ikl', A, X)
print(B.shape) # (3,4,4)

In words, the string ij,jkl->ikl means:

  • A has indices (dimensions) i and j respectively
  • X has indices j, k and l respectively
  • ->ikl multiply A[i,j]*X[j,k,l] and sum over j

Upvotes: 4

Related Questions