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
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.