Умножить матрицы высокого порядка с numpy

Вы можете использовать np.einsum для выполнения операции, поскольку она позволяет очень тщательно контролировать, какие оси умножаются и которые суммируются:

 >>> np.einsum('ijk,ij->ik', ind, dist) array([[ 0.4, 0.4, 0.4, 0.4], [ 3. , 3. , 3. , 3. ], [ 1. , 1. , 1. , 1. ]]) 

Функция умножает записи на первой оси ind с элементами первой оси dist (индекс 'i' ). То же самое для второй оси каждого массива (индекс 'j' ). Вместо того, чтобы возвращать 3D-массив, мы сообщаем einsum суммировать значения вдоль оси 'j' , опуская его из выходных индексов, тем самым возвращая 2D-массив.

np.tensordot сложнее применить к этой проблеме. Он автоматически суммирует продукты осей. Тем не менее, мы хотим два набора продуктов, но суммируем только один из них.

Написание np.tensordot(ind, dist, axes=[1, 1]) (как в ответе, на которое вы ссылаетесь) вычисляет правильные значения для вас, но возвращает 3D-массив с формой (3, 4, 3) . Если вы можете позволить себе стоимость памяти большего массива, вы можете использовать:

 np.tensordot(ind, dist, axes=[1, 1])[0].T 

Это дает правильный результат, но потому, что tensordot создает гораздо больший, чем необходимый массив, einsum представляется лучшим вариантом.

После отличного ответа @ ajcr я хотел бы определить, какой метод является самым быстрым, поэтому я использовал timeit :

 import timeit setup_code = """ import numpy as np i,j,k = (300,200,400) ind = np.ones((i,j,k)) #shape=(3L, 2L, 4L) dist = np.random.rand(i,j) #shape=(3L, 2L) """ basic ="np.array([np.dot(dist[l],ind[l]) for l in xrange(dist.shape[0])])" einsum = "np.einsum('ijk,ij->ik', ind, dist)" tensor= "np.tensordot(ind, dist, axes=[1, 1])[0].T" print "tensor - total time:", min(timeit.repeat(stmt=tensor,setup=setup_code,number=10,repeat=3)) print "basic - total time:", min(timeit.repeat(stmt=basic,setup=setup_code,number=10,repeat=3)) print "einsum - total time:", min(timeit.repeat(stmt=einsum,setup=setup_code,number=10,repeat=3)) 

Неожиданные результаты:

 tensor - total time: 6.59519493952 basic - total time: 0.159871203461 einsum - total time: 0.263569731028 

Таким образом, очевидно, что использование tenordot было неправильным способом сделать это (не говоря уже memory error в более крупных примерах, как указано в @ajcr).

Поскольку этот пример был небольшим, я изменил размер матриц на i,j,k = (3000,200,400) , перевернул порядок, чтобы убедиться, что он не действует, и настроил еще один тест с большим количеством повторений:

 print "einsum - total time:", min(timeit.repeat(stmt=einsum,setup=setup_code,number=50,repeat=3)) print "basic - total time:", min(timeit.repeat(stmt=basic,setup=setup_code,number=50,repeat=3)) 

Результаты были согласованы с первым запуском:

 einsum - total time: 13.3184077671 basic - total time: 8.44810031351 

Однако тестирование другого типа роста размера – i,j,k = (30000,20,40) привело к следующим результатам:

 einsum - total time: 0.325594117768 basic - total time: 0.926416766397 

См. Комментарии для объяснения этих результатов.

Мораль заключается в том, чтобы искать наиболее быстрое решение для конкретной проблемы, пытаться генерировать данные, которые как можно ближе к исходным данным, с точки зрения типов и форм. В моем случае i намного меньше j,k и поэтому я остался с уродливой версией, которая также является самой быстрой в этом случае.