如何在numpy中优雅地构造下列矩阵？

Question

假设我有一个5维矩阵v，现在我想要一个新的矩阵D

D[a, b, n, m, d] = v[a, b, n, n, d]-v[a, b, m, m, d].

我该如何优雅地在numpy中做到这一点？你知道吗

Answer 1

您希望如何更改维度？你可以这样重塑它

import numpy as np

a, b, n, d = 2, 3, 4, 5
v = np.zeros((a, b, n, n, d))
D = v.reshape((a, b, n*n, d))

Answer 2

我发现einsum可以做到这一点：

D = np.einsum('abiic->abic', v)[..., None, :] - np.einsum('abiic->abic', v)[:, :, None, ...]