how to use efficiently rotation matrix in numpy to compute new coordinates along a circle (avoiding for loops)?

Question

I would like improve the script below. The goal of the script is to compute different coordinates corresponding to several angles. This corresponds to rotate one point around x axis and for each step angle retrieve the coordinates.

import numpy as np

# rotating angle
th = np.pi
# starting point (x,y,z coordinates)
ori = [1,1,0]
# how many steps ?
step = 10 

# initializing coordinate matrix
# with the starting point

coords = np.array(ori)

# vector containing angles
angles = np.linspace(0,th,step)

# compute new position for each angle in vector angles
for angle in angles: 
    # adapt Rot matrix with angle
    Rot = np.array([
        [np.cos(angle),-np.sin(angle),0],
        [np.sin(angle),np.cos(angle),0],
        [0, 0, 1]
    ])
    # compute new coordinate
    new_coord = ori @ Rot
    # add new coordinate to coords matrix
    coords = np.append(coords, new_coord)
# change shape of coords for esthetical purpose
coords= np.reshape(coords,(-1,3))
print(coords)

I would like to avoid the loop !

Does anyone have an idea for using only matrix calculations to achieve the same result?

Thank you

Answer 1

Your code, run with %%timeit in an ipython session times as

597 µs ± 1.86 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

The coords is a (11,3) array. The first 2 rows are identical, the ori you first put in it, and the 0 angle rotation.

Replacing the array append with a faster list append produces the same thing, with a modest time improvement

... 
...:   coords = [ori]

...: for angle in angles: 
...:     ...
...:     coords.append(new_coord)
...: # change shape of coords for esthetical purpose
...: coords= np.reshape(coords,(-1,3))

Modest time improvement:

427 µs ± 2.15 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

We can calculate a rot matrix for all angles at once:

def rotmat(angles):
    Rot = np.zeros((len(angles),3,3))
    Rot[:,2,2] = 1
    Rot[:,[0,1],[0,1]] = np.cos(angles)[:,None]
    Rot[:,[0,1],[1,0]] = np.sin(angles)[:,None]*[-1,1]
    return Rot

And apply it with ori@rotmat(angles)

In [64]: %%timeit 
    ...: # rotating angle
    ...: th = np.pi
    ...: # starting point (x,y,z coordinates)
    ...: ori = [1,1,0]
    ...: # how many steps ?
    ...: step = 10 
    ...: 
    ...: # initializing coordinate matrix
    ...: # with the starting point
    ...: 
    ...: # vector containing angles
    ...: angles = np.linspace(0,th,step)
    ...: rot = rotmat(angles)
    ...: coords = ori @ rot
    ...: 
    ...: 
143 µs ± 364 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

This matches your original coords (except it doesn't duplicate the first row):

np.allclose(original_coords[1:],coords)

That's a decent time improvement by using standard numpy 'vectorization' - using whole array methods where possible. numba etc can improve things more.