Numpy: Make Batch Version of Quaternion Multiplication

Question

I converted the following function

def quaternion_multiply(quaternion0, quaternion1): """Return multiplication of two quaternions. >>> q = quaternion_multiply([1, -2, 3, 4], [-5, 6, 7, 8]) >>> numpy.allclose(q, [-44, -14, 48, 28]) True """ x0, y0, z0, w0 = quaternion0 x1, y1, z1, w1 = quaternion1 return numpy.array(( x1*w0 + y1*z0 - z1*y0 + w1*x0, -x1*z0 + y1*w0 + z1*x0 + w1*y0, x1*y0 - y1*x0 + z1*w0 + w1*z0, -x1*x0 - y1*y0 - z1*z0 + w1*w0), dtype=numpy.float64) 

in batch version

 def quat_multiply(self, quaternion0, quaternion1): x0, y0, z0, w0 = np.split(quaternion0, 4, 1) x1, y1, z1, w1 = np.split(quaternion1, 4, 1) result = np.array(( x1*w0 + y1*z0 - z1*y0 + w1*x0, -x1*z0 + y1*w0 + z1*x0 + w1*y0, x1*y0 - y1*x0 + z1*w0 + w1*z0, -x1*x0 - y1*y0 - z1*z0 + w1*w0), dtype=np.float64) return np.transpose(np.squeeze(result)) 

This function processes quaternion1 and quaternion0 with form (?, 4). Now I want the function to be able to handle an arbitrary number of dimensions, such as (?,?, 4). How to do it?

Answer 1

You can get your behavior simply by going axis-=-1 to np.split to split along the last axis.

And since your arrays have this dimensional dimension size 1, and do not fit into a new dimension and then compress it, you can simply combine them again (with the last) axis=-1 :

 def quat_multiply(self, quaternion0, quaternion1): x0, y0, z0, w0 = np.split(quaternion0, 4, axis=-1) x1, y1, z1, w1 = np.split(quaternion1, 4, axis=-1) return np.concatenate( (x1*w0 + y1*z0 - z1*y0 + w1*x0, -x1*z0 + y1*w0 + z1*x0 + w1*y0, x1*y0 - y1*x0 + z1*w0 + w1*z0, -x1*x0 - y1*y0 - z1*z0 + w1*w0), axis=-1) 

Please note that with this approach, you can not only multiply the same forms of quaternion stacks of any number of dimensions:

 >>> a = np.random.rand(6, 5, 4) >>> b = np.random.rand(6, 5, 4) >>> quat_multiply(None, a, b).shape (6, 5, 4) 

But you also get a beautiful broadcast that allows you to multiply the quaternion stack by one without having to bother with the sizes:

 >>> a = np.random.rand(6, 5, 4) >>> b = np.random.rand(4) >>> quat_multiply(None, a, b).shape (6, 5, 4) 

Or, with minimal pile, do all cross-products between two stacks on the same line:

 >>> a = np.random.rand(6, 4) >>> b = np.random.rand(5, 4) >>> quat_multiply(None, a[:, None], b).shape (6, 5, 4) 

Answer 2

You can use np.rollaxis to bring the last axis to the forefront, helping us slice 4 arrays without actually separating them. We perform the required operations and, finally, send the first axis to the end so that the shape of the output array is the same as the inputs. So we would have a solution for generic n-dimensional ndarrays, for example:

 def quat_multiply_ndim(quaternion0, quaternion1): x0, y0, z0, w0 = np.rollaxis(quaternion0, -1, 0) x1, y1, z1, w1 = np.rollaxis(quaternion1, -1, 0) result = np.array(( x1*w0 + y1*z0 - z1*y0 + w1*x0, -x1*z0 + y1*w0 + z1*x0 + w1*y0, x1*y0 - y1*x0 + z1*w0 + w1*z0, -x1*x0 - y1*y0 - z1*z0 + w1*w0), dtype=np.float64) return np.rollaxis(result,0, result.ndim) 

Run Example -

 In [107]: # N-dim arrays ...: a1 = np.random.randint(0,9,(2,3,2,4)) ...: b1 = np.random.randint(0,9,(2,3,2,4)) ...: In [108]: quat_multiply_ndim(a1,b1) # New ndim approach Out[108]: array([[[[ 154., 48., 55., -57.], [ 31., 81., 29., -95.]], [[ 31., 14., 88., 12.], [ 3., 30., 20., -51.]], [[ 104., 61., 102., -39.], [ 0., 14., 14., -56.]]], [[[ -28., 36., 24., -8.], [ 11., 76., -7., -36.]], [[ 54., 3., -2., -19.], [ 52., 62., 15., -55.]], [[ 76., 28., 28., -60.], <--------| [ 14., 54., 13., 5.]]]]) | | In [109]: quat_multiply(a1[1,2],b1[1,2]) # Old 2D approach Out[109]: | array([[ 76., 28., 28., -60.], ------------------| [ 14., 54., 13., 5.]]) 

Answer 3

You are almost there! You just need to be a little careful about how you split and merge your array:

 def quat_multiply(quaternion0, quaternion1): x0, y0, z0, w0 = np.split(quaternion0, 4, axis=-1) x1, y1, z1, w1 = np.split(quaternion1, 4, axis=-1) return np.squeeze(np.stack(( x1*w0 + y1*z0 - z1*y0 + w1*x0, -x1*z0 + y1*w0 + z1*x0 + w1*y0, x1*y0 - y1*x0 + z1*w0 + w1*z0, -x1*x0 - y1*y0 - z1*z0 + w1*w0), axis=-1), axis=-2) 

Here we use axis=-1 both times to split along the last axis and then concatenate backward along the last axis. Finally, we squeeze the second-last axis, as you correctly noted. And show you that it works:

 >>> q0 = np.array([-5, 6, 7, 8]) >>> q1 = np.array([1, -2, 3, 4]) >>> q0 = np.tile(q1, (2, 2, 1)) >>> q0 array([[[-5, 6, 7, 8], [-5, 6, 7, 8]], [[-5, 6, 7, 8], [-5, 6, 7, 8]]]) >>> q1 = np.tile(q2, (2, 2, 1)) >>> q = quat_multiply(q0, q1) array([[[-44, -14, 48, 28], [-44, -14, 48, 28]], [[-44, -14, 48, 28], [-44, -14, 48, 28]]]) >>> q.shape (2, 2, 4) 

Hope you need it! This should work on arbitrary sizes and an arbitrary number of measurements.

Note. np.split does not work in lists. Thus, you can only pass arrays to your new function, as I said above. If you want to be able to transfer lists, you can call instead

  np.split(np.asarray(quaternion0), 4, -1) 

inside your function.

Also, your test case seems to be wrong. I think you swapped the positions of quaternion0 and quaternion1 : I swapped them back when testing q0 and q1 .