cross products with einsums

Question

Я пытаюсь вычислить перекрестные произведения многих векторных пар 3x1 как можно быстрее. Это

n = 10000 
a = np.random.rand(n, 3) 
b = np.random.rand(n, 3) 
numpy.cross(a, b) 

дает правильный ответ, но мотивировано this answer to a similar question, я думал, что einsum поймают меня куда-то. Я обнаружил, что оба

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->ai', eijk, a, b) 
np.einsum('iak,ak->ai', np.einsum('ijk,aj->iak', eijk, a), b) 

вычислить декартово произведение, но их эффективность вызывает разочарование: Оба метода работают намного хуже, чем np.cross:

%timeit np.cross(a, b) 
1000 loops, best of 3: 628 µs per loop 

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

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

Любой идеи о том, как улучшить einsum s?

Answer 1

Количество операции умножения на einsum() является более cross(), и в новой версии NumPy, cross() не создает много временных массивов. Таким образом, einsum() не может быть быстрее, чем cross().

Вот старый код креста:

x = a[1]*b[2] - a[2]*b[1] 
y = a[2]*b[0] - a[0]*b[2] 
z = a[0]*b[1] - a[1]*b[0] 

Вот новый код креста:

multiply(a1, b2, out=cp0) 
tmp = array(a2 * b1) 
cp0 -= tmp 
multiply(a2, b0, out=cp1) 
multiply(a0, b2, out=tmp) 
cp1 -= tmp 
multiply(a0, b1, out=cp2) 
multiply(a1, b0, out=tmp) 
cp2 -= tmp 

ускорив его, вам нужно Cython или Numba.

Answer 2

Вы можете ввести матрицу-умножение, используя np.tensordot потерять одно из измерений на первый уровень, а затем использовать np.einsum потерять другое измерение, например, так -

np.einsum('aik,ak->ai',np.tensordot(a,eijk,axes=([1],[1])),b) 

В качестве альтернативы, мы можем выполнить транслируемые поэлементно умножений между a и b использованием np.einsum, а затем потерять два размера в один проход с np.tensordot, как так -

np.tensordot(np.einsum('aj,ak->ajk', a, b),eijk,axes=([1,2],[1,2])) 

Мы могли бы выполнить элементарные умножения, вводя новые оси тоже с чем-то вроде a[...,None]*b[:,None], но, похоже, это замедляет его.

---

Хотя, они показывают хорошее улучшение по сравнению с предлагаемыми np.einsum подходов только на основе, но не бить np.cross.

Продолжительность испытания -

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->ai', eijk, a, b) 
    ...: %timeit np.einsum('iak,ak->ai', np.einsum('ijk,aj->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->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->ajk',a,b),eijk,axes=([1,2],[1,2])) 
1000 loops, best of 3: 882 µs per loop