计算每个箱子的最大数量

Question

我正在努力执行一个概念上简单的算法使用优化numpy矢量化操作。在下面的代码中，我有一个data，它是一个带有一堆值的数组，而coords，其中条目i包含与data[i]对应的三维空间坐标。我想填充数组max_data，这样条目max_data[i,j,k]就是data所有条目的最大值，这样coords对应的条目就在[ [i,i+1], [j,j+1], [k,k+1] ]范围内。下面是生成数据并实现算法的示例代码。你知道吗

有没有什么方法可以使用numpy矢量化操作来加快速度？我在使用ndata~1e9的阵列上运行了一个版本，这需要很长时间。我并不反对使用其他python库。你知道吗

import numpy as np
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 

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