Search code examples
pythonnumpyvectorvectorizationquaternions

Applying quaternion rotation to a vector time series


I have a time series of 3D vectors in a Python numpy array similar to the following:

array([[-0.062, -0.024,  1.   ],
       [-0.071, -0.03 ,  0.98 ],
       [-0.08 , -0.035,  0.991],
       [-0.083, -0.035,  0.98 ],
       [-0.083, -0.035,  0.977],
       [-0.082, -0.035,  0.993],
       [-0.08 , -0.034,  1.006],
       [-0.081, -0.032,  1.008],
       .......

I want to rotate each vector around a specified axis through a specified angle theta. I have been using quaternions to achieve this for one vector as found here in henneray's answer.

v1 = np.array ([1, -2, 0])
axis = np.array([-4, -2,  3])
theta = 1.5

rot_axis = np.insert(axis, 0, 0, axis=0)
axis_angle = (theta*0.5) * rot_axis/np.linalg.norm(rot_axis)
vec = quat.quaternion(*v1)
qlog = quat.quaternion(*axis_angle)
q = np.exp(qlog)
v_prime = q * vec * np.conjugate(q)
v_prime_vec = v_prime.imag

My question is, what is the fastest way to apply the same rotation to each vector in v1?

You cannot create a quaternion from v1 if v1 contains a 2D array of vectors, so I could use a loop to rotate each array element in turn; however, in henneray's answer in the link above, it is mentioned that the quaternions could be applied to 'appropriately vectorised numpy arrays'. Does anyone has any suggestions on how this could be implemented?

(A side question: if my theta and axis variables were arrays of equal length to v1, could the same method also be used to rotate each vector in v1 through a corresponding rotation?)


Solution

  • It is necessary to first convert the [x,y,z] Cartesian vectors into 4-vectors with the first component equal to zero [0,x,y,z]. Then you can cast this to a quaternion array to do vectorised calculations.

    This function below takes an array of Cartesian vectors and rotates them about a single rotation axis. You will need to make sure the norm of this axis is equal to your rotation angle theta.

    def rotate_vectors(vecs, axis):
        """
        Rotate a list of 3D [x,y,z] vectors about corresponding 3D axis
        [x,y,z] with norm equal to the rotation angle in radians
    
        Parameters
        ----------
        vectors : numpy.ndarray with shape [n,3]
            list of [x,y,z] cartesian vector coordinates
        axis : numpy.ndarray with shape [3]
            [x,y,z] axis to rotate corresponding vectors about
        """
        # Make an 4 x n array of zeros
        vecs4 = np.zeros([vecs.shape[0],vecs.shape[1]+1])
        # Fill the imaginary i, j, k components with x, y, z values, leaving the real part w=0
        vecs4[:,1:] = vecs
        # Convert to quaternion array
        vecsq = quat.as_quat_array(vecs4)
    
        # Make a rotation quaternion
        qrot = quat.from_rotation_vector(axis)
        # Rotate vectors
        vecsq_rotated = qrot * vecsq * qrot.conjugate()
        # Cast quaternion array to float and return only imaginary components (ignore real part)
        return quat.as_float_array(vecsq_rotated)[:,1:]
    

    As a bonus, this function takes an array of rotation axes to rotate each vector by the corresponding axes.

    def rotate_vectors_each(vecs, axes):
        """
        Rotate a list of 3D [x,y,z] vectors about corresponding 3D axes
        [x,y,z] with norm equal to the rotation angle in radians
    
        Parameters
        ----------
        vectors : numpy.ndarray with shape [n,3]
            list of [x,y,z] cartesian vector coordinates
        axes : numpy.ndarray with shape [n,3]
            axes to rotate corresponding vectors about
            n = pulse shape time domain
            3 = [x,y,z]
        """
        # Make an 4 x n array of zeros
        vecs4 = np.zeros([vecs.shape[0],vecs.shape[1]+1])
        # Fill the imaginary i, j, k components with x, y, z values, leaving the real part w=0
        vecs4[:,1:] = vecs
        # Convert to quaternion array
        vecsq = quat.as_quat_array(vecs4)
    
        # Make an 4 x n array of zeros
        rots4 = np.zeros([rots.shape[0],rots.shape[1]+1])
        # Fill the imaginary i, j, k components with x, y, z values, leaving the real part w=0
        rots4[:,1:] = rots
        # Convert to quaternion array and take exponential
        qrots = np.exp(quat.as_quat_array(0.5 * rots4))
    
        # Rotate vectors
        vecsq_rotated = qrots * vecsq * qrots.conjugate()
    
        return quat.as_float_array(vecsq_rotated)[:,1:]
    

    Note that with so many conversions between axis angle and quaternion representation, this will give you little performance improvement over rotation matrix algebra. Quaternions really only benefit when you are rotating a vector through many sequential rotations, whereby you can stack the quaternion multiplication.