How to speed up rotation operations in Python?

Question

I have an N-body simulation script that works with thousand of particles. It outputs a final 2D projection of the position of the particles, and I want to run it multiple times with different rotation angles around a given axis, to be able to see the final 2D result from different perspectives. To accomplish this, I added a little code at the beginning:

# Euler-Rodrigues formula definition for arbitrary 3D rotation
def rotation_matrix(axis,angle):
    axis = axis/math.sqrt(dot(axis,axis))
    a = math.cos(angle/2)
    b,c,d = -axis*math.sin(angle/2)
    return array([[a*a+b*b-c*c-d*d, 2*(b*c-a*d), 2*(b*d+a*c)],
                  [2*(b*c+a*d), a*a+c*c-b*b-d*d, 2*(c*d-a*b)],
                  [2*(b*d-a*c), 2*(c*d+a*b), a*a+d*d-b*b-c*c]])

And then:

# 3D rotation
angle = float(sys.argv[7])*math.pi/180.0
axis = array([0,0,1])
xAr,yAr,zAr=[],[],[]           #arrays for rotated particles
for xi,yi,zi in zip(xA,yA,zA): #cycle to rotate every particle
    ci = array([xi,yi,zi])
    cr = dot(rotation_matrix(axis,angle),ci)
    xAr.append(cr[0])
    yAr.append(cr[1])
    zAr.append(cr[2])
xA,yA,zA = array(xAr), array(yAr), array(zAr)

The core part of the script starts after this. In summary, originally the scripts does this:

• reads model > runs simulation > outputs a 2D projection

and I added my part so that now:

• reads model > rotates every particle in 3D > runs simulation > outputs a 2D projection

but I found that the rotation process takes too much time. Is there a way or a different approach in Python to speed up this? (if it helps, I only want the final 2D projection rotated along a given axis).

Answer 1

You need to vectorize the operation, i.e. run it on the whole set of particles instead of looping over them. Assuming that xA etc. are 1-d arrays, that's something like:

particles = np.vstack([xA, xB, xC])
rot = rotation_matrix(axis, angle)
rotated = np.dot(rot, particles)
xA, yA, zA = rotated

This should give you several orders of magnitude speedup. On a more general note, you're constructing the same rotation matrix in each iteration of the loop. That's wasteful.

Answer 2

What about puting rotation_matrix(axis,angle) out of tight loop