如何使用numpy apply_along_axis

问题描述

我有一个可以完全矢量化的问题，但是我没有足够的空间，所以我正在尝试使用numpy的apply_along_axis（）实现一半的解决方案。

（注意：这是一个说明问题的玩具示例。换句话说，我不是在寻找执行此处功能的numpy或scipy函数-它不是真正的函数，一个简单的例子。）

我想做的是找出一种方法来访问每次迭代传递的轴的索引。

假设我们采用4 x 4矩阵：

    M = np.array(([0,1,1],[1,0],[0,1]))
    M 
   
   array([[0,1]])

并且想要计算每列相对于其他每一列的成对按位逻辑和，但是为了节省（大量）时间，我们仅计算i> j的列，其中j> i的索引（这样我们最终带有三角形矩阵）。

在熊猫中，我可以使用apply（）很容易地做到这一点，但是对于我的目的来说太慢了。

我知道scikit-learn中有成对函数，但是请假定这些函数不符合我的目的（我的函数比这个玩具更复杂）

如果我要使用numpy的apply_long_axis（），则只能算出如何比较所有i，j和j，i，而不是前面所述的较小问题。

这是我的解决方案：

def intersections_np(col,M):
    col = col[:,np.newaxis]
    intersection = (M & col).sum(0)
    return(intersection)

result_np = np.apply_along_axis(intersections_np,arr = M,axis = 0,M = M)
result_np

array([[2,3,2],2,3]],dtype=int32)

但是我真正想做的是：

def intersections_np(col,np.newaxis]
    start_index = <index_of_current_column> + 1
    other_cols = M[:,start_index:]
    intersection = (other_cols & col).sum(0)
    <possible padding of the array with nans here>
    return(intersection)

result_np = np.apply_along_axis(intersections_np,M = M)

并返回：

result_np

array([[nan,nan,nan],nan]],dtype=int32)

有人知道这样的事情可以做吗？

谢谢

解决方法

让我们做些时间。

您的基本apply：

In [142]: timeit np.apply_along_axis(intersections_np,arr = M,axis = 0,M = M)                    
158 µs ± 3.97 µs per loop (mean ± std. dev. of 7 runs,10000 loops each)

等效迭代（从技术上讲，可能需要转置，结果是对称的，所以没关系）：

In [143]: timeit np.array([intersections_np(M[:,i],M) for i in range(M.shape[1])])                   
65.4 µs ± 1.93 µs per loop (mean ± std. dev. of 7 runs,10000 loops each)

和@jfahne建议：

In [144]: %%timeit  
     ...: np.reshape(np.array([(M.T[i] & M.T[j]).sum(0) if j>i else 0 \ 
     ...: for i in range(len(M.T)) for j in range(len(M.T))]),(M.T).shape).T 
     ...:  
     ...:                                                                                            
95.2 µs ± 2.99 µs per loop (mean ± std. dev. of 7 runs,10000 loops each)

请注意，apply比普通迭代要慢。这与我过去的测试是一致的。 apply仅在数组为3d或更大时才有帮助，并且迭代是“丑陋”的双嵌套数组。那里比较漂亮，虽然还不算快。这是一种便利工具，而不是速度工具。

完全“矢量化”的解决方案（具有numpy广播等）

In [148]: (M[:,:,None] & M[:,None,:]).sum(0)                                                         
Out[148]: 
array([[2,1,1],[1,3,2],2,3]])
In [149]: timeit (M[:,:]).sum(0)                                                  
14.9 µs ± 182 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)

它确实制作了一个中间（4,4,4）数组，并且没有避免重复，但是因为在Python级别上没有迭代，所以它非常快。试图将计算限制在较低（或较高）三角形上通常是不值得的。

但是，如果您确实想要较低的三角形和速度，请考虑使用numba。对于迭代问题，它可能很快（但要付出一定的灵活性）。

---

这是您的相交版本仅限于下三角形

In [159]: def foo(M): 
     ...:     m = M.shape[0] 
     ...:     res = np.full((m,m),np.nan) 
     ...:     for i in range(m-1): 
     ...:         temp = (M[:,i,(i+1):]).sum(0) 
     ...:         res[-temp.shape[0]:,i] = temp 
     ...:     return res 
     ...:      
     ...:                                                                                            
In [160]: foo(M)                                                                                     
Out[160]: 
array([[nan,nan,nan],[ 1.,0.,1.,2.,nan]])
In [161]: timeit foo(M)                                                                              
59.3 µs ± 2.42 µs per loop (mean ± std. dev. of 7 runs,10000 loops each)

与我的[143]基本相同的计时-&步骤中的计算较少，但是索引较多，因此速度变化很小。

,

这有一个相当不错的pythonic解决方案：

np.reshape(np.array([(M.T[i] & M.T[j]).sum(0) if j>i else 0 \
for i in range(len(M.T)) for j in range(len(M.T))]),(M.T).shape).T

使用M.T来访问列。生成的矢量将重塑为与转置数组相同的形状。然后将阵列转置回原始阵列形状，并产生所需的输出。