Multiplying a matrix with array of vectors in NumPy

Question

I'm trying to geometrically rotate an array of vectors in NumPy. First I generate the coordinate vectors of a grid.

width = 128
height = 128

x_axis = np.linspace(-1, 1, width)
y_axis = np.linspace(-1, 1, height)

x, y = np.meshgrid(x_axis, y_axis)
z = np.full((width, height), 0)
vectors = np.stack((x, y, z), axis=2)

So 'vectors' has the shape (128, 128, 3)

I already prepared the rotation matrix with a, b and c as the rotation angles along the axes.

rotation_matrix = np.array([
    [np.cos(b) * np.cos(c),
     - np.cos(b) * np.sin(c),
     np.sin(b)],

    [np.sin(a) * np.sin(b) * np.cos(c) + np.cos(a) * np.sin(c),
     - np.sin(a) * np.sin(b) * np.sin(c) + np.cos(a) * np.cos(c),
     - np.sin(a) * np.cos(b)],

    [- np.cos(a) * np.sin(b) * np.cos(c) + np.sin(a) * np.sin(c),
     np.cos(a) * np.sin(b) * np.sin(c) + np.sin(a) * np.cos(c),
     np.cos(a) * np.cos(b)]
])

Now I want every vector of the array to be matrix multiplied with 'rotation_matrix' like

vector_rotated = rotation_matrix @ vector

so the resulting array also should have the shape (128, 128, 3). I have some problems with handling this 3 dimensional array. Matmul is only capable of handling 2d arrays. Is there any elegant way in NumPy for this use case or do I have to use a for loop to solve this issue?

Thanks a lot and have a nice day!

Answer 1

There are a few different ways to solve this problem.

Option 1:

The most straightforward is to reshape the array vectors so that it has shape (3, 128 * 128) , then call the builtin np.dot function, and reshape the result back to your desired shape.

(Note that the (128, 128) part of the array's shape is not really relevant to the rotation; it's an interpretation that you probably want to make your problem clearer, but makes no difference to the linear transformation you want to apply. Said another way, you are rotating 3-vectors. There are 128 * 128 == 16384 of them, they just happen to be organized into a 3D array like above.)

This approach would look like:

>>> v = vectors.reshape(-1, 3).T
>>> np.dot(rotation_matrix, v).shape
(3, 16384)
>>> rotated = np.dot(rotation_matrix, v).T.reshape(vectors.shape)
>>> rotated.shape == vectors.shape
True

Option 2:

Another method that does not involve any reshaping is to use NumPy's Einstein summation . Einstein summation is very flexible, and takes a while to understand, but its power justifies its complexity. In its simplest form, you "label" the axes that you want to multiply together. Axes that are omitted are "contracted", meaning a sum across that axis is computed. For your case, it would be:

>>> np.einsum('ij,klj->kli', rotation_matrix, vectors).shape
(128, 128, 3)
>>> np.allclose(rotated, np.einsum('ij,klj->kli', rotation_matrix_vectors))
True

Here's a quick explanation of the indexing. We are labeling the axes of the rotation matrix i and j , and the axes of the vectors k , l , and j . The repeated j means those are the axes multiplied together. This is equivalent to right-multiplying the reshaped array above with the rotation matrix (ie, it's a rotation).

The output axes are labeled kli . This means we're preserving the k and l axes of the vectors. Since j is not in the output labels, there is a summation across that axis. Instead we have the axis i , hence the final shape of (128, 128, 3) . You can see above that the dot-product method and the einsum method agree.

It can take a while to wrap your head around Einstein summation, but it is super awesome and powerful. I highly recommend learning more about it, especially if this sort of linear algebra is a common problem for you.