计算一系列 numpy ndarrays 的外积

Question

我有一个 3D 点列表 p 存储在形状为 (N, 3) 的 ndarray 中。我想计算每个 3d 点的外积：

N = int(1e4)
p = np.random.random((N, 3))
result = np.zeros((N, 3, 3))
for i in range(N):
    result[i, :, :] = np.outer(p[i, :], p[i, :])

有没有一种方法可以在不使用任何 python 级循环的情况下计算这个外积？问题是 np.outer 不支持任何类似 axis 的论点。

Answer 1

你至少可以使用 apply_along_axis:

result = np.apply_along_axis(lambda point: np.outer(point, point), 1, p)

然而，令人惊讶的是，这实际上 比您的方法慢：

In [ ]: %%timeit N = int(1e4); p = np.random.random((N, 3))
   ...: result = np.apply_along_axis(lambda point: np.outer(point, point), 1, p)
61.5 ms ± 1.84 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [ ]: %%timeit N = int(1e4); p = np.random.random((N, 3))
   ...: result = np.zeros((N, 3, 3))
   ...: for i in range(N):
   ...:     result[i, :, :] = np.outer(p[i, :], p[i, :])
46 ms ± 709 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

您可以使用广播：

p[..., None] * p[:, None, :]

此语法在第一项的末尾（使其成为 Nx3x1）和第二项的中间（使其成为 Nx1x3）插入一个轴。然后将这些广播并产生 Nx3x3 结果。

Answer 3

比我以前的解决方案更好的解决方案是使用 np.einsum:

np.einsum('...i,...j', p, p)

比广播方式还要快：

In [ ]: %timeit p[..., None] * p[:, None, :]
514 µs ± 4.23 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [ ]: %timeit np.einsum('...i,...j', p, p)
169 µs ± 1.75 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

至于它是如何工作的我不是很确定，我只是乱搞einsum直到我得到我想要的答案：

In [ ]: np.all(np.einsum('...i,...j', p, p) == p[..., None] * p[:, None, :])
Out[ ]: True

计算一系列 numpy ndarrays 的外积

Calculating the outer product for a sequence of numpy ndarrays

python

performance

numpy

matrix

scientific-computing