带有任意运算符的 numpy 点或 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

您可以使用 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 2

numpy github 上有一个老问题，要求对 einsum 进行泛化，以允许使用其他函数。当前版本只是实现了产品的总和。据我所知，没有人接过那个项目。

几年前我修补了 einsum，修复了 '...' 符号的处理。所以我很清楚它是如何实现的；并可能调整我的 Python/cython 模拟器来添加此功能。实际的einsum代码写在c.

我的猜测是，如果您不喜欢 Divakar 的方法，则必须使用 cython 编写您自己的版本。

Answer 3

您可以将 Divakar 的回答中的广播与 numexpr 一起使用：

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

1*() 是 a workaround. 我已经确认这不会分配一个大的临时数组。

带有任意运算符的 numpy 点或 einsum

numpy dot or einsum with arbitrary operator

numpy

numpy-ufunc