CODEWITHSUNDEEP
pythonnumpy
Patrick Rinker
Patrick Rinker

Reputation: 315

Multiplying matrix A times multi-dimensional matrix "matrix-wise?"

I have a 3d numpy array u, shape (k, m, n), and I am trying to compute a new array uprod, shape (k, m, n), such that uprod[j] = np.dot(A, u[j]), where A is a fixed matrix that doesn't depend on j at all. I can easily do this writing a loop over the inner-most index, but is there a faster/better way to do so?

Upvotes: 2

Views: 71

Answers (1)

user2357112
user2357112

Reputation: 280837

np.einsum can do the job:

result = numpy.einsum('ij,kjl->kil', A, u)

You can also make it broadcast, so if X and Y are thought of as arrays of 2D matrices, the following call will perform an appropriately broadcasted dot:

result = numpy.einsum('...ij,...jk->...ik', X, Y)

For example, if X has shape (3, 4, 5, 6) and Y has shape (4, 6, 5), then result[1, 2] would be an array of shape (5, 5) equal to X[1, 2].dot(Y[2]).

You can also do this with dot. A.dot(u) produces a result array with A.dot(u)[i, j, k] == A[i, :].dot(u[j, :, k). You want a result array with result[i, j, k] == A[j, :].dot(u[i, :, k]); you can get this with rollaxis or transpose with an axes argument.

result = numpy.rollaxis(A.dot(u), 1)

Where einsum makes broadcasting easy, dot for high-dimensional input is kind of like an outer product. With the same X and Y as before, if you set

result = numpy.rollaxis(X.dot(Y), axis=X.ndim-2, start=X.ndim+Y.ndim-3)

then result[1, 2, 3] would be an array of shape (5, 5) equal to X[1, 2].dot(Y[3]).

Upvotes: 3

Related Questions