Вычисление максимума numpy ndarray в каждой корзине

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] ]. Пример кода, который генерирует данные и реализует алгоритм, приведен ниже.

Можно ли как-нибудь ускорить это, используя простые векторизованные операции? Я запускаю версию этого на массивах с 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

Использование 2-го самого быстрого решения из 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