NumPy中的外部产品:矢量化六个嵌套循环 [英] Exterior product in NumPy : Vectorizing six nested loops

问题描述

在研究论文中，作者介绍了两个(3 * 3)矩阵A和B之间的外部乘积，得出C:

In a research paper, the author introduces an exterior product between two (3*3) matrices A and B, resulting in C:

C(i, j) = sum(k=1..3, l=1..3, m=1..3, n=1..3) eps(i,k,l)*eps(j,m,n)*A(k,m)*B(l,n)

其中eps(a, b, c)是 Levi-Civita符号.

我想知道如何在Numpy中向量化这样的数学运算符，而不是天真地实现6个嵌套循环(对于i, j, k, l, m, n).

I am wondering how to vectorize such a mathematical operator in Numpy instead of implementing 6 nested loops (for i, j, k, l, m, n) naively.

推荐答案

它看起来像是一个纯粹基于和约的问题，不需要在输入之间保持任何轴对齐.因此，我建议使用 np.tensordot .

It looks like a purely sum-reduction based problem without the requirement of keeping any axis aligned between the inputs. So, I would suggest matrix-multiplication based solution for tensors using np.tensordot.

因此，一种解决方案可以分三步实施-

Thus, one solution could be implemented in three steps -

# Matrix-multiplication between first eps and A.
# Thus losing second axis from eps and first from A : k 
parte1 = np.tensordot(eps,A,axes=((1),(0)))

# Matrix-multiplication between second eps and B.
# Thus losing third axis from eps and second from B : n
parte2 = np.tensordot(eps,B,axes=((2),(1)))

# Finally, we are left with two products : ilm & jml. 
# We need to lose lm and ml from these inputs respectively to get ij. 
# So, we need to lose last two dims from the products, but flipped .
out = np.tensordot(parte1,parte2,axes=((1,2),(2,1)))

运行时测试

方法-

def einsum_based1(eps, A, B):  # @unutbu's soln1
    return np.einsum('ikl,jmn,km,ln->ij', eps, eps, A, B)

def einsum_based2(eps, A, B):  # @unutbu's soln2
    return np.einsum('ilm,jml->ij',
              np.einsum('ikl,km->ilm', eps, A), 
              np.einsum('jmn,ln->jml', eps, B))

def tensordot_based(eps, A, B):
    parte1 = np.tensordot(eps,A,axes=((1),(0)))
    parte2 = np.tensordot(eps,B,axes=((2),(1)))
    return np.tensordot(parte1,parte2,axes=((1,2),(2,1)))

时间-

In [5]: # Setup inputs
   ...: N = 20
   ...: eps = np.random.rand(N,N,N)
   ...: A = np.random.rand(N,N)
   ...: B = np.random.rand(N,N)
   ...: 

In [6]: %timeit einsum_based1(eps, A, B)
1 loops, best of 3: 773 ms per loop

In [7]: %timeit einsum_based2(eps, A, B)
1000 loops, best of 3: 972 µs per loop

In [8]: %timeit tensordot_based(eps, A, B)
1000 loops, best of 3: 214 µs per loop

更大的数据集-

In [12]: # Setup inputs
    ...: N = 100
    ...: eps = np.random.rand(N,N,N)
    ...: A = np.random.rand(N,N)
    ...: B = np.random.rand(N,N)
    ...: 

In [13]: %timeit einsum_based2(eps, A, B)
1 loops, best of 3: 856 ms per loop

In [14]: %timeit tensordot_based(eps, A, B)
10 loops, best of 3: 49.2 ms per loop

这篇关于NumPy中的外部产品:矢量化六个嵌套循环的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！