numpy:列式点积

Question

给定一个2D numpy数组,我需要计算每列的点积,并将结果存储在一维数组中.以下作品:

In [45]: A = np.array([[1,2,3,4],[5,6,7,8]])

In [46]: np.array([np.dot(A[:,i], A[:,i]) for i in xrange(A.shape[1])])
Out[46]: array([26, 40, 58, 80])

有一种简单的方法可以避免Python循环吗？以上几乎不是世界末日,但如果有numpy这个原因,我想用它.

编辑在实践中,矩阵有许多行和相对较少的列.因此,我并不过分热衷于创建大于的临时数组O(A.shape[1]).我也无法修改A到位.

Answer 1

怎么样:

>>> A = np.array([[1,2,3,4],[5,6,7,8]])
>>> (A*A).sum(axis=0)
array([26, 40, 58, 80])

编辑:嗯,好吧,你不想要中间的大型物体.也许:

>>> from numpy.core.umath_tests import inner1d
>>> A = np.array([[1,2,3,4],[5,6,7,8]])
>>> inner1d(A.T, A.T)
array([26, 40, 58, 80])

反正似乎有点快.这应该在幕后做你想要的,因为AT是一个视图(它没有自己的副本,IIUC),而inner1d 似乎以它需要的方式循环.

非常新的更新:另一种选择是使用np.einsum:

>>> A = np.array([[1,2,3,4],[5,6,7,8]])
>>> np.einsum('ij,ij->j', A, A)
array([26, 40, 58, 80])
>>> timeit np.einsum('ij,ij->j', A, A)
100000 loops, best of 3: 3.65 us per loop
>>> timeit inner1d(A.T, A.T)
100000 loops, best of 3: 5.02 us per loop
>>> A = np.random.randint(0, 100, (2, 100000))
>>> timeit np.einsum('ij,ij->j', A, A)
1000 loops, best of 3: 363 us per loop
>>> timeit inner1d(A.T, A.T)
1000 loops, best of 3: 848 us per loop
>>> (np.einsum('ij,ij->j', A, A) == inner1d(A.T, A.T)).all()
True