二维阵列上的矢量化NumPy三积

Question

我正在尝试对下面名为p的N x N数组上的以下三重乘积操作进行矢量化：

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

只需将这些循环迭代器作为字符串符号进行移植，我们就可以得到一个基于^{}的解决方案，如下所示-

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

基本上是一个扩展相关的问题（因为我们没有减少任何轴），我们也可以简单地使用^{}，通过在不同的地方引入新的轴来允许扩展，就像这样-

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