Question

我正在努力使用优化的numpy向量化运算来执行概念上简单的算法。在下面的代码中，我有data，这是一个带有一堆值的数组，并且其中coords的项i包含与data[i]相对应的3D空间坐标。我想填充数组max_data，使条目max_data[i,j,k]是data的所有条目的最大值，以使coords的相应条目落入{{1} }。下面是生成数据并实现算法的示例代码。

是否有任何方法可以使用numpy向量化操作来加快速度？我正在使用ndata〜1e9在阵列上运行此版本，这需要花很长时间。我不反对使用其他python库。

[ [i,i+1], [j,j+1], [k,k+1] ]

Answer 1

使用https://stackoverflow.com/a/55226663/7207392中第二快的解决方案可以使>30x加速。如果您愿意使用pythran，则可以使用更快的解决方案。

import numpy as np
from scipy import sparse
import time

shape = ( 20, 30, 40 )

ndata = int( 1e6 )

data = np.random.normal(  loc = 10, scale = 5, size = ndata ) 

coords = np.vstack( [ np.random.uniform( 0, shape[i], ndata )
                      for i in range( len( shape ) ) ] ).T

max_data = np.zeros( shape ) 

start = time.time()

for i in range( len( data ) ) :

    # shortcut to find bin indices when the bins are
    # [ range( shape[i] ) for i in range( len( shape ) ) ]

    bin_indices = tuple( coords[i].astype( int ) )  

    max_data[ bin_indices ] = max( max_data[ bin_indices ], data[ i ] )

elapsed = time.time() - start

print( 'elapsed: %.3e' % elapsed )  # 2.98 seconds on my computer 


start = time.time()

bin_indices = np.ravel_multi_index(coords.astype(int).T, shape)
aux = sparse.csr_matrix((data, bin_indices, np.arange(data.size+1)),
                        (data.size, np.prod(shape))).tocsc()
cut = aux.indptr.searchsorted(data.size)
max_data_pp = np.empty(shape)
max_data_pp.ravel()[:cut] = np.maximum.reduceat(aux.data, aux.indptr[:cut])

CLIPAT = 0

max_data_pp.ravel()[aux.indptr[:-1]==aux.indptr[1:]] = CLIPAT
max_data_pp[max_data_pp < CLIPAT] = CLIPAT

elapsed = time.time() - start

print( 'elapsed: %.3e' % elapsed )  # 2.98 seconds on my computer 


assert np.all(max_data == max_data_pp)

样品运行：

elapsed: 2.417e+00
elapsed: 6.387e-02

计算每个bin中的numpy ndarray的最大值

1 个答案: