Python 计算每个箱子中的最大numpy ndarray数

我正在努力使用优化的numpy矢量化操作来执行一个概念上简单的算法。在下面的代码中，我有

数据

，它是一个具有一组值的数组，并且

坐标

其中的条目

I

包含与

数据[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 

---

使用中第二快的解决方案，我的速度提高了30倍。如果您愿意使用

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)

---

针对任何面临类似问题的人更新：我发现

scipy.stats.binned\u statistic\u dd

，它正是我想要的，但有趣的是，它比问题中的原始代码慢10倍多。谢谢。顺便说一句，

CLIPAT

表示什么？@jacob它是用来模拟代码的行为的，因为您使用

0进行初始化

会忽略负面数据。

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