在二维数组上向量化 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

只需将这些循环迭代器移植为字符串符号，我们就会有一个基于 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]

在二维数组上向量化 NumPy 三重积

vectorize NumPy triple product on 2D array

python

arrays

numpy

vectorization

numpy-einsum