How to multiply the same matrix with each row of a tensor?

I wanted to find coordinates of a cube after two (X and Y) rotations without loops.

The cube coordinates are stored as vertices:

cube_vertices = np.array([[0, 0, 0], 
                          [1, 0, 0], 
                          [1, 1, 0], 
                          [0, 1, 0], 
                          [0, 0, 1], 
                          [1, 0, 1], 
                          [0, 1, 1], 
                          [1, 1, 1]])

And i mostly work with edges:

cube_edges = np.array([[cube_vertices[0], cube_vertices[1]],
                       [cube_vertices[0], cube_vertices[3]],
                       [cube_vertices[1], cube_vertices[2]],
                       [cube_vertices[3], cube_vertices[2]],
                       [cube_vertices[1], cube_vertices[5]],
                       [cube_vertices[2], cube_vertices[7]],
                       [cube_vertices[3], cube_vertices[6]],
                       [cube_vertices[0], cube_vertices[4]],
                       [cube_vertices[4], cube_vertices[5]],
                       [cube_vertices[4], cube_vertices[6]],
                       [cube_vertices[6], cube_vertices[7]],
                       [cube_vertices[5], cube_vertices[7]]])

that I store to find intersections of these edges with plane. Rotator is given as rotations repeated 12 times for each coordinate of cube_edges, so i could multiply along 0-th axis of size 12 (number of edges) and 1-st axis of size 2 (two ends of an edge) for all the coordinates:

phi, theta = 0.1, 0.1
Rx = np.tile(np.array([[1, 0, 0], 
                       [0, np.cos(theta), -np.sin(theta)], 
                       [0, np.sin(theta), np.cos(theta)]]), (12, 2, 1, 1))

Ry = np.tile(np.array([[np.cos(phi), 0, np.sin(phi)], 
                       [0, 1, 0], 
                       [-np.sin(phi), 0, np.cos(phi)]]), (12, 2, 1, 1))
Rxy = Rx * Ry

The shape of cube_edges is (12, 2, 3) and shape of Rxy is (12, 2, 3, 3). The multiplication should act along first two dimensions (12, 2) for all 3x3 single rotations per point coordinates of size (3). I can't figure out how to utilize np.tensordot or someting like that for this purpose:

rotated_edges = np.tensordot(Rxy, cube_edges)

The resulting shape of rotated_edges should be the same as cube_edges (12, 2, 3).

Thanks!

Solution

The following seems to do what you're attempting (using the same array cube_edges):

phi, theta = 0.1, 0.1
Rx = np.array([[1, 0, 0], 
               [0, np.cos(theta), -np.sin(theta)], 
               [0, np.sin(theta), np.cos(theta)]])

Ry = np.array([[np.cos(phi), 0, np.sin(phi)], 
               [0, 1, 0], 
               [-np.sin(phi), 0, np.cos(phi)]])
Rxy = Rx @ Ry
rotated_edges = np.dot(cube_edges,Rxy.T)

Here's a little self-contained demo you can play with:

import numpy as np
import matplotlib.pyplot as plt

## setup
cube_vertices = np.array([[0, 0, 0], 
                          [1, 0, 0], 
                          [1, 1, 0], 
                          [0, 1, 0], 
                          [0, 0, 1], 
                          [1, 0, 1], 
                          [0, 1, 1], 
                          [1, 1, 1]])
edges = [
         [0, 1], [0, 3], [0, 4],
         [1, 2], [1, 5],
         [2, 7], [2, 3],
         [3, 6],
         [4, 5], [4, 6],
         [5, 7], [6, 7]
]
cube_edges = np.array([[cube_vertices[i], cube_vertices[j]] for i,j in edges])

## rotate
phi, theta = 0.1, 0.1
Rx = np.array([[1, 0, 0], 
               [0, np.cos(theta), -np.sin(theta)], 
               [0, np.sin(theta), np.cos(theta)]])

Ry = np.array([[np.cos(phi), 0, np.sin(phi)], 
               [0, 1, 0], 
               [-np.sin(phi), 0, np.cos(phi)]])
Rxy = Rx @ Ry
rotated_edges = np.dot(cube_edges,Rxy.T)

## plot
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
for pair in cube_edges:
    ax.plot3D(*pair.T)
plt.show()

fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
for pair in rotated_edges:
    ax.plot3D(*pair.T)
plt.show()

The resulting graphs I get: