将高阶矩阵与numpy相乘

Question

我创造了这个玩具问题,反映了我更大的问题:

import numpy as np

ind = np.ones((3,2,4)) # shape=(3L, 2L, 4L)
dist = np.array([[0.1,0.3],[1,2],[0,1]]) # shape=(3L, 2L)

ans = np.array([np.dot(dist[i],ind[i]) for i in xrange(dist.shape[0])]) # shape=(3L, 4L)
print ans

""" prints:
   [[ 0.4  0.4  0.4  0.4]
    [ 3.   3.   3.   3. ]
    [ 1.   1.   1.   1. ]]
"""

我想尽快做到这一点,所以使用numpy的函数来计算ans应该是最好的方法,因为这个操作很重,我的矩阵很大.

我看到这篇文章,但形状不同,我无法理解axes我应该用于这个问题.但是,我确信,数字应该有答案.有什么建议？

编辑:我接受了@ ajcr的答案,但也请阅读我自己的答案,它可能会帮助别人......

Answer 1

您可以使用它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')的第二个轴.我们告诉einsum不是返回3D数组,而是'j'通过从输出下标中省略它来沿轴对它们求和,从而返回一个2D数组.

np.tensordot更难以应用于这个问题.它自动汇总轴的产品.但是,我们希望2套产品,但总结只有一个人.

写入np.tensordot(ind, dist, axes=[1, 1])(如您链接到的答案中)为您计算正确的值,但返回带有形状的3D数组(3, 4, 3).如果你能承受更大阵列的内存成本,你可以使用:

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

这会给你正确的结果,但是因为首先tensordot创建一个比必要的大得多的数组,einsum似乎是一个更好的选择.

Answer 2

按照@ 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

所以显然使用tensordot是错误的做法(更不用说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我小得多,所以我留在丑陋的版本,这也是这种情况下最快的.