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

我想使用类似于

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

的每一行中的相应元素数。我希望能够执行以下等效操作（注意：下面的输出是虚构的，但值是我希望看到的）：

有什么方法可以做到这一点吗？

您可以使用这种方法来解决匹配计数问题-

(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))

---

numpy github上有一个老问题，要求对

einsum

进行泛化，以允许使用其他函数。当前版本只实现了一系列产品。据我所知，还没有人接手那个项目

几年前，我修补了

einsum

，修复了“…”符号的处理。所以我对它是如何实施的有一个很好的想法；并且可能会修改我的Python/cython仿真器以添加此功能。实际的

einsum

代码是用

c

编写的

---

我的猜测是，如果您不喜欢Divakar的方法，您必须使用

cython

编写自己的版本，您可以使用Divakar答案中的广播以及：

---

1*（）

是我已经确认这不会分配一个大的临时数组。

这确实产生了预期的结果。一个缺点是

（a[：，None]==b）

和

np.equal（a[：，None]，b）

的临时存储。在这个例子中，中间数组的形状是

（3,4,2）

，但实际上，最后一个维度可能很大，这就是我希望避免它的原因。@bogatron实际用例中

a

和

b

的典型形状是什么？每个数组的第一个维度大约是10**5。第二个维度（求和发生的维度）可能是64或128。@bogatron那么匹配的布尔数组

a[：，None]==b

，因为是布尔数组，所以比int/float数组效率更高。即使这样，如果您遇到内存问题，我们也可以在cols中运行一个循环，执行类似于a[：，i，None]==b[：，i]的操作，并在每次迭代中累积。除此之外，我已经没有想法了。@Divakar-numexpr；-）但是您必须使用

1*（a[：，None]==b）

，因为数组总是只有

0

s和

1

s吗？在我的特定情况下，是的。但是我对更一般的情况也很好奇。我认为

numexpr

不支持切片，很高兴知道它支持！这只是

local_dict

参数，切片仍然在Python中发生：-）我认为这将最有效，因为它避免了临时数组存储。Upvote@Divakar是elegent numpy唯一的解决方案。接受的答案符合我的需要，但如果看到

einsum

支持此功能，那就太好了。我发现，对于许多操作，它的运行速度比其他解决方案快得多。

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

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