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编写你自己的版本。

numpy dot或einsum with any operator

3 个答案: