python: как эффективно создать список трехмерных точек

Question

У меня есть длинный список «точек» в 3D, у каждой из которых есть [x,y,z,scalar,scalar,...]

Я хотел бы «связать» точки в 3D (в x / y / z) или повторно -организовать в новый «массив массивов» с измерениями [n_bins_x, n_bins_y, n_bins_z], каждый элемент которого является подмассивом. Мне не нужно просто «считать» их, как с помощью гистограммы, скорее, я бы хотел получить массив точек для каждого бина, упакованный в трехмерный массив.

В моем реальном случае использования используется O (1M) баллов * O (2K) временных шагов, каждый из которых должен быть связан, поэтому необходима высокая производительность.

Вот пример:

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import timeit

# ----- generate pts

size_x, size_y, size_z = 5.0, 3.0, 1.0
n_pts = 100000

x = np.random.uniform(0.0, size_x, n_pts )
y = np.random.uniform(0.0, size_y, n_pts )
z = np.random.uniform(0.0, size_z, n_pts )
u = np.random.uniform(0.0, 1.0,    n_pts )
v = np.random.uniform(0.0, 1.0,    n_pts )
w = np.random.uniform(0.0, 1.0,    n_pts )
data = np.vstack((x,y,z,u,v,w)).T

# ----- define bin bounds

n_bins_x, n_bins_y, n_bins_z = 10, 10, 10
n_bins_total = n_bins_x*n_bins_y*n_bins_z
x_bin_bounds = np.linspace(0.0, size_x, n_bins_x+1, dtype=np.float64)
y_bin_bounds = np.linspace(0.0, size_y, n_bins_y+1, dtype=np.float64)
z_bin_bounds = np.linspace(0.0, size_z, n_bins_z+1, dtype=np.float64)

# ----- np.digitize to get 'bin index' for each particle in each dim
# --> this basically creates 3 new scalars which are integer-valued 'flags'
#      for the sorting process later

bin_inds_x = np.digitize(data[:,0], x_bin_bounds)-1
bin_inds_y = np.digitize(data[:,1], y_bin_bounds)-1
bin_inds_z = np.digitize(data[:,2], z_bin_bounds)-1

# ----- re-organize total list into a 3D array of sub-arrays

data_resorted = np.empty(shape=(n_bins_x, n_bins_y, n_bins_z), dtype=object)

inds_all = np.array(range(n_pts), dtype=np.int32)

# ----- method 1

start_time = timeit.default_timer()
counter=1; particle_counter=0
for i in range(n_bins_x):
    for j in range(n_bins_y):
        for k in range(n_bins_z):
            #print('%i/%i'%(counter,n_bins_total)); counter+=1
            a = np.where(bin_inds_x==i) ## if in the current searched-for x-bin
            b = np.where(bin_inds_y==j) ## if in the current searched-for y-bin
            c = np.where(bin_inds_z==k) ## if in the current searched-for z-bin
            d = np.intersect1d(a, b) ## if in x AND y bin being searched for
            e = np.intersect1d(c, d) ## if in x AND y AND z bin being searched for
            inds = inds_all[e] ## take the particles meeting those criteria
            count = len(inds)
            if (count > 0):
                particle_counter += count
                data_resorted[i,j,k] = data[inds,:] ## assign matched points to 're-organized' array

end_time = timeit.default_timer() - start_time
print('method 1 time: %0.2f[s]'%end_time)

if (particle_counter == n_pts):
    print('all pts binned %i %i'%(particle_counter, n_pts))
else:
    print('pts lost! %i %i'%(particle_counter, n_pts))

# ----- method 2

start_time = timeit.default_timer()
counter=1; particle_counter=0
for i in range(n_bins_x):
    inds = np.copy(inds_all)
    inds = inds[np.where(bin_inds_x[inds]==i)] ## index 'mask' for this x-bin
    inds_copy_i = np.copy(inds) ## matches in (x)... copy for re-use in nested (y,z) searches
    for j in range(n_bins_y):                  
        inds = np.copy(inds_copy_i)
        inds = inds[np.where(bin_inds_y[inds]==j)] ## index 'mask' for this y-bin
        inds_copy_j = np.copy(inds) ## matches in (x,y)... copy for re-use in nested (z) searches
        for k in range(n_bins_z):
            #print('%i/%i'%(counter,n_bins_total)); counter+=1
            inds = np.copy(inds_copy_j)
            inds = inds[np.where(bin_inds_z[inds]==k)] ## index 'mask' for this z-bin
            count = len(inds)
            if (count > 0):
                particle_counter += count
                data_resorted[i,j,k] = data[inds,:] ## assign matched points to 're-organized' array
end_time = timeit.default_timer() - start_time
print('method 2 time: %0.2f[s]'%end_time)

if (particle_counter == n_pts):
    print('all pts binned %i %i'%(particle_counter, n_pts))
else:
    print('pts lost! %i %i'%(particle_counter, n_pts))

Вывод:

method 1 time: 1.79[s]
all pts binned 100000 100000
method 2 time: 0.04[s]
all pts binned 100000 100000

«метод 2» демонстрирует большое преимущество в производительности по сравнению с «методом 1», но у меня все еще возникает ощущение, что я не правильно использую numpy для этой задачи, которая может быть выполнена в более «векторизованный» способ внутри numpy или использование другой функции numpy.

Существует ли лучший, более эффективный способ «складирования» частиц в 3D в python / numpy?