带有任意运算符的 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 编写的。

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