кросс-продукты с айнсумами

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

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

дает правильный ответ, но, мотивированный этим ответом на аналогичный вопрос, я подумал, что 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

Любые идеи о том, как улучшить einsums?


python numpy performance numpy-einsum cross-product
person Nico Schlömer    schedule 23.09.2016    source источник

Ответы (2)


arrow_upward
4
arrow_downward

Количество операций умножения 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.

person HYRY    schedule 23.09.2016
comment
И поскольку они не включают циклы, улучшение cython может быть незначительным. При таком выражении cross является скорее алгебраической операцией, чем операцией массива. - person hpaulj; 23.09.2016

arrow_upward
4
arrow_downward

Вы можете ввести матричное умножение, используя 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
person Divakar    schedule 23.09.2016