numpy dot или einsum с произвольным оператором

Question

Я хотел бы использовать что-то вроде np.dot или (желательно) np.einsum для эффективной работы с ними, но с альтернативным ufunc вместо np.multiply. Например, рассмотрим эти два массива:

>>> a 
array([[0, 1], 
     [1, 1], 
     [1, 0]]) 
>>> b 
array([[0, 0], 
     [1, 0], 
     [1, 0], 
     [0, 0]]) 

Теперь предположим, что я хочу, чтобы подсчитать количество элементов в каждой строке a равной соответствующих элементов в каждой строке b. Я хотел бы быть в состоянии сделать эквивалент следующего (примечание: ниже выход изготовлен, но значения то, что я ожидал увидеть):

>>> np.dot(a, b.T, ufunc=np.equal) 
array([[1, 0, 0, 1], 
     [0, 1, 1, 0], 
     [1, 2, 2, 1]]) 

Есть ли способ сделать это?

Answer 1

Вы можете использовать вещанию от ответа Divakar вместе с numexpr:

numexpr.evaluate('sum(1*(a == b), axis=2)', {'a': a[:,None]}) 

1*() является a workaround. Я подтвердил это не выделяет большой временный массив.

Answer 2

Вы можете использовать broadcasting для такого подсчета матч проблема -

(a[:,None] == b).sum(2) 

Пример выполнения -

In [36]: a 
Out[36]: 
array([[0, 1], 
     [1, 1], 
     [1, 0]]) 

In [37]: b 
Out[37]: 
array([[0, 0], 
     [1, 0], 
     [1, 0], 
     [0, 0]]) 

In [38]: (a[:,None] == b).sum(2) 
Out[38]: 
array([[1, 0, 0, 1], 
     [0, 1, 1, 0], 
     [1, 2, 2, 1]]) 

Если вы действительно хотите использовать np.einsum и np.equal, вот способ формовать ранее подход к дать нам желаемый результат -

np.einsum('ijk->ij',np.equal(a[:,None],b).astype(int)) 

Answer 3

Существует старая проблема в numpy github с просьбой обобщить einsum, что позволит использовать другие функции. Текущая версия просто реализует сумму продуктов. Насколько я знаю, никто этого не предпринял.

Несколько лет назад я исправил einsum, установив обработку обозначения «...». Поэтому я имею хорошее представление о том, как это реализовано; и, возможно, адаптировать мой эмулятор Python/cython для добавления этой функции. Фактический код einsum написан на c.

Я предполагаю, что если вам не нравится подход Дивакара, вам придется написать свою собственную версию с помощью cython.