cross products with einsums
<p>I'm trying to compute the cross-products of many 3x1 vector pairs as fast as possible. This</p>



<pre><code>n = 10000
a = np.random.rand(n, 3)
b = np.random.rand(n, 3)
numpy.cross(a, b)
</code></pre>

<p>gives the correct answer, but motivated by <a href="http://stackoverflow.com/a/20910319/353337">this answer to a similar question</a>, I thought that <code>einsum</code> would get me somewhere. I found that both</p>



<pre><code>eijk = np.zeros((3, 3, 3))
eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1

np.einsum('ijk,aj,ak-&gt;ai', eijk, a, b)
np.einsum('iak,ak-&gt;ai', np.einsum('ijk,aj-&gt;iak', eijk, a), b)
</code></pre>

<p>compute the cross product, but their performance is disappointing: Both methods perform much worse than <code>np.cross</code>:</p>

<pre><code>%timeit np.cross(a, b)
1000 loops, best of 3: 628 Âµs per loop
</code></pre>



<pre><code>%timeit np.einsum('ijk,aj,ak-&gt;ai', eijk, a, b)
100 loops, best of 3: 9.02 ms per loop
</code></pre>



<pre><code>%timeit np.einsum('iak,ak-&gt;ai', np.einsum('ijk,aj-&gt;iak', eijk, a), b)
100 loops, best of 3: 10.6 ms per loop
</code></pre>

<p>Any ideas of how to improve the <code>einsum</code>s?</p>

<p>You can bring in matrix-multiplication using <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.tensordot.html" rel="nofollow"><code>np.tensordot</code></a> to lose one of the dimensions at the first level and then use <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html" rel="nofollow"><code>np.einsum</code></a> to lose the other dimension, like so -</p>

<pre><code>np.einsum('aik,ak-&gt;ai',np.tensordot(a,eijk,axes=([1],[1])),b)
</code></pre>

<p>Alternatively, we can perform broadcasted elementwise multiplications between <code>a</code> and <code>b</code> using <code>np.einsum</code> and then lose the two dimensions in one go with <code>np.tensordot</code>, like so -</p>

<pre><code>np.tensordot(np.einsum('aj,ak-&gt;ajk', a, b),eijk,axes=([1,2],[1,2]))
</code></pre>

<p>We could have performed the elementwise multiplications by introducing new axes too with something like <code>a[...,None]*b[:,None]</code>, but it seems to slow it down.</p>

<hr>

<p>Though, these show good improvement over the proposed <code>np.einsum</code> only based approaches, but fail to beat <code>np.cross</code>.</p>

<p>Runtime test -</p>

<pre><code>In [26]: # Setup input arrays
    ...: n = 10000
    ...: a = np.random.rand(n, 3)
    ...: b = np.random.rand(n, 3)
    ...: 

In [27]: # Time already posted approaches
    ...: %timeit np.cross(a, b)
    ...: %timeit np.einsum('ijk,aj,ak-&gt;ai', eijk, a, b)
    ...: %timeit np.einsum('iak,ak-&gt;ai', np.einsum('ijk,aj-&gt;iak', eijk, a), b)
    ...: 
1000 loops, best of 3: 298 Âµs per loop
100 loops, best of 3: 5.29 ms per loop
100 loops, best of 3: 9 ms per loop

In [28]: %timeit np.einsum('aik,ak-&gt;ai',np.tensordot(a,eijk,axes=([1],[1])),b)
1000 loops, best of 3: 838 Âµs per loop

In [30]: %timeit np.tensordot(np.einsum('aj,ak-&gt;ajk',a,b),eijk,axes=([1,2],[1,2]))
1000 loops, best of 3: 882 Âµs per loop
</code></pre>

python performance numpy cross-product numpy-einsum