python - 在二维数组上矢量化 NumPy 三重乘积

Question

我正在尝试在下面N x N调用的数组上对以下三重乘积运算进行矢量化p：

for j in range(len(p)):
    for k in range(len(p)):
        for l in range(len(p)):
            h[j, k, l] = p[j, k] * p[k, l] * p[l, j] - p[j, l] * p[l, k] * p[k, j]

我认为numpy.einsum应该在这里使用，尽管我实际上并没有对重复的索引进行求和，但我无法确定它。想法？

Answer 1

只需将这些循环迭代器作为字符串表示法移植，我们就会有一个einsum像这样的基于解决方案 -

h = np.einsum('jk,kl,lj->jkl',p,p,p) - np.einsum('jl,lk,kj->jkl',p,p,p)

基本上是一个与扩展相关的问题（因为我们没有减少任何轴），我们也可以NumPy broadcasting通过在不同位置引入新轴来简单地使用None/np.newaxis以允许扩展，就像这样 -

h = p[...,None]*p*p[:,None,:].T - p[:,None,:]*p.T*p.T[...,None]