¿Cómo puedo vectorizar este triple lazo en 2d matrices en numpy?

Question

¿Puedo eliminar todos los bucles de Python en este cálculo:

result[i,j,k] = (x[i] * y[j] * z[k]).sum() 

donde x[i], y[j], z[k] son vectores de longitud N y x, y, z tienen primeras dimensiones con una longitud A, B, C S.T. la salida tiene forma (A,B,C) y cada elemento es la suma de un producto triple (elemento-sabio).

Lo puedo bajar de 3 a 1 bucles (código a continuación), pero estoy atascado tratando de eliminar el último bucle.

Si es necesario, podría hacer A=B=C (mediante una pequeña cantidad de relleno).

# Example with 3 loops, 2 loops, 1 loop (testing omitted) 

N = 100 # more like 100k in real problem 
A = 2 # more like 20 in real problem 
B = 3 # more like 20 in real problem 
C = 4 # more like 20 in real problem 

import numpy 
x = numpy.random.rand(A, N) 
y = numpy.random.rand(B, N) 
z = numpy.random.rand(C, N) 

# outputs of each variant 
result_slow = numpy.empty((A,B,C)) 
result_vec_C = numpy.empty((A,B,C)) 
result_vec_CB = numpy.empty((A,B,C)) 

# 3 nested loops 
for i in range(A): 
    for j in range(B): 
     for k in range(C): 
      result_slow[i,j,k] = (x[i] * y[j] * z[k]).sum() 

# vectorize loop over C (2 nested loops) 
for i in range(A): 
    for j in range(B): 
     result_vec_C[i,j,:] = (x[i] * y[j] * z).sum(axis=1) 

# vectorize one C and B (one loop) 
for i in range(A): 
    result_vec_CB[i,:,:] = numpy.dot(x[i] * y, z.transpose()) 

numpy.testing.assert_almost_equal(result_slow, result_vec_C) 
numpy.testing.assert_almost_equal(result_slow, result_vec_CB) 

Answer 1

Si está utilizando numpy> 1.6, no es el impresionante np.einsum función:

np.einsum('im,jm,km->ijk',x,y,z) 

lo que equivale a sus versiones en bucle. No estoy seguro de cómo va a ser justo en cuanto a la eficiencia una vez que se llega al tamaño de sus matrices en el problema real (en realidad, recibo una segfault en mi máquina, cuando me muevo a esos tamaños). La otra solución que a menudo prefiero para este tipo de problemas es volver a escribir el método usando cython.

Answer 2

El uso de einsum tiene mucho sentido en su caso; pero puedes hacerlo fácilmente con la mano. El truco consiste en hacer que las matrices sean genéricamente estables unas contra otras. Eso significa remodelarlos para que cada conjunto varíe independientemente a lo largo de su propio eje. Luego multiplíquelos juntos, dejando que numpy se encargue de la transmisión; y luego sumar a lo largo del último eje (más a la derecha).

>>> x = numpy.arange(2 * 4).reshape(2, 4) 
>>> y = numpy.arange(3 * 4).reshape(3, 4) 
>>> z = numpy.arange(4 * 4).reshape(4, 4) 
>>> (x.reshape(2, 1, 1, 4) * 
... y.reshape(1, 3, 1, 4) * 
... z.reshape(1, 1, 4, 4)).sum(axis=3) 
array([[[ 36, 92, 148, 204], 
     [ 92, 244, 396, 548], 
     [ 148, 396, 644, 892]], 

     [[ 92, 244, 396, 548], 
     [ 244, 748, 1252, 1756], 
     [ 396, 1252, 2108, 2964]]]) 

Puede hacer esto un poco más generalizada mediante el uso de la notación de la rebanada, el valor newaxis (que es igual a None, por lo que el siguiente sería trabajar con None también), y el hecho de que sum acepta valores de los ejes negativos (con -1 que denota el último, -2 que indica el penúltimo, y así sucesivamente). De esta forma, no es necesario que conozca la forma original de las matrices; siempre que sus últimos ejes sean compatibles, transmitirá los tres primeros juntos:

>>> (x[:, numpy.newaxis, numpy.newaxis, :] * 
... y[numpy.newaxis, :, numpy.newaxis, :] * 
... z[numpy.newaxis, numpy.newaxis, :, :]).sum(axis=-1) 
array([[[ 36, 92, 148, 204], 
     [ 92, 244, 396, 548], 
     [ 148, 396, 644, 892]], 

     [[ 92, 244, 396, 548], 
     [ 244, 748, 1252, 1756], 
     [ 396, 1252, 2108, 2964]]])