Exterior product in NumPy : Vectorizing six nested loops

Question

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)

where eps(a, b, c) is the Levi-Civita symbol.

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.

Answer 1

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)))

Runtime test

Approaches -

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)))

Timings -

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

Bigger dataset -

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