在numpy中获取matrix <matrix的结果的最快方法是什么?

时间:2018-06-20 07:52:34

标签: python performance numpy vectorization

假设我有一个维度为(M,A)的矩阵M_1和一个维度为(M,B)的矩阵M_2M_1 < M_2的结果应为维度(M,B,A)的矩阵,其中M1中的每一行都与M_2的相应行中的每个元素进行比较,并给出每个比较的布尔向量(或1,0-向量)。

例如,如果我有一个矩阵

M1 = [[1,2,3]
      [3,4,5]]

M2 = [[1,2],
      [3,4]]

result should be [[[False, False, False],
                   [True, False, False]],
                  [[False, False, False], 
                   [True, False, False]]]

当前,我正在使用for循环,当我不得不多次重复此操作(耗时数月)时,该循环非常慢。希望有一种矢量化的方法可以做到这一点。如果没有,我还能做什么?

我正在查看M_1为(500,3000000)和M_2为(500,500),并重复了大约10000次。

1 个答案:

答案 0 :(得分:3)

对于NumPy数组,使用None/np.newaxis扩展暗角,以使第一个轴对齐,而第二个轴被 spread 扩展,使它们可以按元素方式进行比较。最后,利用broadcasting进行比较,以获得向量化解决方案-

M1[:,None,:] < M2[:,:,None]

样品运行-

In [19]: M1
Out[19]: 
array([[1, 2, 3],
       [3, 4, 5]])

In [20]: M2
Out[20]: 
array([[1, 2],
       [3, 4]])

In [21]: M1[:,None,:] < M2[:,:,None]
Out[21]: 
array([[[False, False, False],
        [ True, False, False]],

       [[False, False, False],
        [ True, False, False]]])

对于列表作为输入,请使用numpy.expand_dims,然后进行比较-

In [42]: M1 = [[1,2,3],
    ...:       [3,4,5]]
    ...: 
    ...: M2 = [[1,2],
    ...:       [3,4]]

In [43]: np.expand_dims(M1, axis=1) < np.expand_dims(M2, axis=2)
Out[43]: 
array([[[False, False, False],
        [ True, False, False]],

       [[False, False, False],
        [ True, False, False]]])

进一步增强

利用multi-core with numexpr module处理大数据时进一步提高-

In [44]: import numexpr as ne

In [52]: M1 = np.random.randint(0,9,(500, 30000))

In [53]: M2 = np.random.randint(0,9,(500, 500))

In [55]: %timeit M1[:,None,:] < M2[:,:,None]
1 loop, best of 3: 3.32 s per loop

In [56]: %timeit ne.evaluate('M1e<M2e',{'M1e':M1[:,None,:],'M2e':M2[:,:,None]})
1 loop, best of 3: 1.53 s per loop
相关问题