Numpy: rendre la version batched de la multiplication de quaternion

Question

I a transformé la fonction suivante

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) 

à une version batched

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

Cette fonction gère quaternion1 et quaternion0 avec la forme (4?). Maintenant, je veux que la fonction peut gérer un nombre arbitraire de dimensions, telles que (?,?, 4). Comment faire ça?

Answer 1

Vous pouvez obtenir le comportement que vous recherchez en passant simplement axis-=-1 à np.split pour diviser le dernier axe.

Et puisque vos tableaux ont cette taille ennuyeux 1 dimension de fuite, plutôt que l'empilement le long d'une nouvelle dimension, puis presser que l'on plus loin, vous pouvez simplement les concaténer, encore une fois le long de la (dernière) 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) 

Notez que, avec cette approche, vous pouvez non seulement multiplier les piles de quaternions forme identique d'un nombre quelconque de dimensions:

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

Mais vous obtenez également la belle diffusion qui permet à IEmultiplier une pile de quaternions avec un seul sans avoir à jouer avec les dimensions:

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

Ou avec tripoter minimum faire tous les produits croisés entre deux piles en une seule ligne:

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

Answer 2

Vous pouvez utiliser np.rollaxis pour amener le dernier axe à l'avant, nous aidant à découper les 4 tableaux sans les diviser. Nous effectuons les opérations requises et finalement envoyer retour le premier axe à la fin pour garder la forme du tableau de sortie même que les entrées. Ainsi, nous aurions une solution pour ndarrays n dimensions génériques, comme si -

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

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

Vous y êtes presque! Vous avez juste besoin d'être un peu prudent sur la façon dont vous partagez et concaténer votre tableau:

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) 

Ici, nous utilisons axis=-1 les deux fois de diviser le long du dernier axe, puis concaténer le long du dernier axe . Enfin, nous pressons l'avant-dernier axe, comme vous l'avez remarqué. Et pour vous montrer que cela fonctionne:

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

Espérons que c'est ce dont vous aviez besoin! Cela devrait fonctionner sur des dimensions arbitraires et un nombre arbitraire de dimensions.

Remarque:np.split semble ne pas fonctionner dans les listes. Vous ne pouvez donc passer des tableaux qu'à votre nouvelle fonction, comme je l'ai fait plus haut. Si vous voulez être en mesure de passer des listes, vous pouvez appeler à la place

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

dans votre fonction.

De même, votre scénario de test semble incorrect. Je pense que vous avez échangé les positions quaternion0 et quaternion1: Je les ai inversés ci-dessus tout en testant q0 et q1.