改变输入数组维度的广播函数答案

【问题标题】：Broadcast function that changes dimension of the input array改变输入数组维度的广播函数
【发布时间】：2016-05-11 18:09:39
【问题描述】：

给定一些函数f 接受1D 数组并给出2D 数组，是否可以有效地将其应用于NxM 数组A 的每一行？

更具体地说，我想为NxM 数组A 的每一行应用np.triu，然后连接所有结果。我可以通过

B = np.dstack(map(np.triu, A))

给出MxMxN 矩阵。但是，对于较大的 N，这不是很有效。不幸的是，这里不能使用函数 np.apply_along_axis，因为 f 改变了维度。

了解 NumPy 在高效广播方面的强大功能后，我几乎可以肯定存在更好的解决方案来解决我的问题。

【问题讨论】：

apply_along_axis 使用 Python 迭代。它只是一个包装器，可以更轻松地沿多维数组的一个轴进行迭代（对于某些程序员而言）。

标签： arrays numpy multidimensional-array array-broadcasting

【解决方案1】：

这是使用broadcasting 的矢量化方法 -

Bout = A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])

运行时测试和输出验证 -

In [319]: A = np.random.randint(0,20,(400,100))

In [320]: %timeit np.dstack(map(np.triu, A))
10 loops, best of 3: 69.9 ms per loop

In [321]: %timeit A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])
10 loops, best of 3: 24.8 ms per loop

In [322]: B = np.dstack(map(np.triu, A))

In [323]: Bout = A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])

In [324]: np.allclose(B,Bout)
Out[324]: True

【讨论】：

或者，如果 dstack 不是故意的，则布局更自然：res=a[:,None,:]*tri(a.shape[1],dtype=bool).T。 res[0] 是 triu(a[0])。
@Divakar 太好了！谢谢。你使用 bool 类型是因为它需要更少的内存还是因为 int 数组与 bool 的乘法更快？
@cheyp 更少的内存是很大的动力。虽然不确定布尔值的性能，但由于通常性能与内存有很好的相关性，我认为这应该是要走的路！ :)
@B.M.谢谢！这也有效，我只是想用面具本身广播，所以没有玩扩展A。感谢您的投入！
还有一个建议：在我的系统上，使用dtype=A.dtype 而不是bool，速度要快20%。