Quel est le moyen le plus rapide de mapper les noms de groupe du tableau numpy aux index?

Je travaille avec pointcloud 3D de Lidar. Les points sont donnés par un tableau numpy qui ressemble à ceci:

points = np.array([[61651921, 416326074, 39805], [61605255, 416360555, 41124], [61664810, 416313743, 39900], [61664837, 416313749, 39910], [61674456, 416316663, 39503], [61651933, 416326074, 39802], [61679969, 416318049, 39500], [61674494, 416316677, 39508], [61651908, 416326079, 39800], [61651908, 416326087, 39802], [61664845, 416313738, 39913], [61674480, 416316668, 39503], [61679996, 416318047, 39510], [61605290, 416360572, 41118], [61605270, 416360565, 41122], [61683939, 416313004, 41052], [61683936, 416313033, 41060], [61679976, 416318044, 39509], [61605279, 416360555, 41109], [61664837, 416313739, 39915], [61674487, 416316666, 39505], [61679961, 416318035, 39503], [61683943, 416313004, 41054], [61683930, 416313042, 41059]])

Je voudrais garder mes données groupées en cubes de taille de 50*50*50sorte que chaque cube conserve un indice hashable et indices numpy de mon pointsqu'il contient . Pour obtenir le fractionnement, j'attribue les cubes = points \\ 50sorties à:

cubes = np.array([[1233038, 8326521, 796], [1232105, 8327211, 822], [1233296, 8326274, 798], [1233296, 8326274, 798], [1233489, 8326333, 790], [1233038, 8326521, 796], [1233599, 8326360, 790], [1233489, 8326333, 790], [1233038, 8326521, 796], [1233038, 8326521, 796], [1233296, 8326274, 798], [1233489, 8326333, 790], [1233599, 8326360, 790], [1232105, 8327211, 822], [1232105, 8327211, 822], [1233678, 8326260, 821], [1233678, 8326260, 821], [1233599, 8326360, 790], [1232105, 8327211, 822], [1233296, 8326274, 798], [1233489, 8326333, 790], [1233599, 8326360, 790], [1233678, 8326260, 821], [1233678, 8326260, 821]])

Ma sortie souhaitée ressemble à ceci:

{(1232105, 8327211, 822): [1, 13, 14, 18]), 
(1233038, 8326521, 796): [0, 5, 8, 9], 
(1233296, 8326274, 798): [2, 3, 10, 19], 
(1233489, 8326333, 790): [4, 7, 11, 20], 
(1233599, 8326360, 790): [6, 12, 17, 21], 
(1233678, 8326260, 821): [15, 16, 22, 23]}

Mon vrai pointcloud contient jusqu'à quelques centaines de millions de points 3D. Quel est le moyen le plus rapide pour effectuer ce type de regroupement?

J'ai essayé une majorité de solutions différentes. Voici une comparaison de la consommation de temps en supposant que la taille des points est d'environ 20 millions et la taille des cubes distincts est d'environ 1 million:

Pandas [tuple (elem) -> np.array (dtype = int64)]

import pandas as pd
print(pd.DataFrame(cubes).groupby([0,1,2]).indices)
#takes 9sec

Defauldict [elem.tobytes () ou tuple -> list]

#thanks @abc:
result = defaultdict(list)
for idx, elem in enumerate(cubes):
    result[elem.tobytes()].append(idx) # takes 20.5sec
    # result[elem[0], elem[1], elem[2]].append(idx) #takes 27sec
    # result[tuple(elem)].append(idx) # takes 50sec

numpy_indexed [int -> np.array]

# thanks @Eelco Hoogendoorn for his library
values = npi.group_by(cubes).split(np.arange(len(cubes)))
result = dict(enumerate(values))
# takes 9.8sec

Pandas + réduction de la dimensionnalité [int -> np.array (dtype = int64)]

# thanks @Divakar for showing numexpr library:
import numexpr as ne
def dimensionality_reduction(cubes):
    #cubes = cubes - np.min(cubes, axis=0) #in case some coords are negative 
    cubes = cubes.astype(np.int64)
    s0, s1 = cubes[:,0].max()+1, cubes[:,1].max()+1
    d = {'s0':s0,'s1':s1,'c0':cubes[:,0],'c1':cubes[:,1],'c2':cubes[:,2]}
    c1D = ne.evaluate('c0+c1*s0+c2*s0*s1',d)
    return c1D
cubes = dimensionality_reduction(cubes)
result = pd.DataFrame(cubes).groupby([0]).indices
# takes 2.5 seconds

Il est possible de télécharger le cubes.npzfichier ici et d'utiliser une commande

cubes = np.load('cubes.npz')['array']

pour vérifier le temps de performance.

Nombre constant d'indices par groupe

Approche n ° 1

Nous pouvons effectuer dimensionality-reductionpour réduire cubesà un tableau 1D. Ceci est basé sur un mappage des données de cubes données sur une grille n-dim pour calculer les équivalents d'indice linéaire, discuté en détail here. Ensuite, en fonction de l'unicité de ces indices linéaires, nous pouvons séparer les groupes uniques et leurs indices correspondants. Par conséquent, en suivant ces stratégies, nous aurions une solution, comme ceci -

N = 4 # number of indices per group
c1D = np.ravel_multi_index(cubes.T, cubes.max(0)+1)
sidx = c1D.argsort()
indices = sidx.reshape(-1,N)
unq_groups = cubes[indices[:,0]]

# If you need in a zipped dictionary format
out = dict(zip(map(tuple,unq_groups), indices))

Alternative n ° 1: si les valeurs entières dans cubessont trop grandes, nous pourrions vouloir faire en dimensionality-reductionsorte que les dimensions avec une extension plus courte soient choisies comme axes principaux. Par conséquent, pour ces cas, nous pouvons modifier l'étape de réduction pour obtenir c1D, comme ceci -

s1,s2 = cubes[:,:2].max(0)+1
s = np.r_[s2,1,s1*s2]
c1D = cubes.dot(s)

Approche n ° 2

Ensuite, nous pouvons utiliser Cython-powered kd-treepour une recherche rapide du plus proche voisin pour obtenir les indices voisins les plus proches et donc résoudre notre cas comme ça -

from scipy.spatial import cKDTree

idx = cKDTree(cubes).query(cubes, k=N)[1] # N = 4 as discussed earlier
I = idx[:,0].argsort().reshape(-1,N)[:,0]
unq_groups,indices = cubes[I],idx[I]

Cas générique: nombre variable d'indices par groupe

Nous allons étendre la méthode basée sur argsort avec un certain fractionnement pour obtenir notre sortie souhaitée, comme ceci -

c1D = np.ravel_multi_index(cubes.T, cubes.max(0)+1)

sidx = c1D.argsort()
c1Ds = c1D[sidx]
split_idx = np.flatnonzero(np.r_[True,c1Ds[:-1]!=c1Ds[1:],True])
grps = cubes[sidx[split_idx[:-1]]]

indices = [sidx[i:j] for (i,j) in zip(split_idx[:-1],split_idx[1:])]
# If needed as dict o/p
out = dict(zip(map(tuple,grps), indices))

Utilisation de versions 1D de groupes de cubesclés as

Nous allons étendre la méthode listée précédemment avec les groupes de cubesclés as pour simplifier le processus de création de dictionnaire et aussi le rendre efficace avec, comme ceci -

def numpy1(cubes):
    c1D = np.ravel_multi_index(cubes.T, cubes.max(0)+1)        
    sidx = c1D.argsort()
    c1Ds = c1D[sidx]
    mask = np.r_[True,c1Ds[:-1]!=c1Ds[1:],True]
    split_idx = np.flatnonzero(mask)
    indices = [sidx[i:j] for (i,j) in zip(split_idx[:-1],split_idx[1:])]
    out = dict(zip(c1Ds[mask[:-1]],indices))
    return out

Ensuite, nous utiliserons numba package pour itérer et arriver à la sortie finale du dictionnaire lavable. Pour aller avec, il y aurait deux solutions - L'une qui obtient les clés et les valeurs séparément en utilisant numbaet l'appel principal sera compressé et converti en dict, tandis que l'autre créera un numba-supportedtype de dict et donc aucun travail supplémentaire requis par la fonction d'appel principale .

Ainsi, nous aurions une première numbasolution:

from numba import  njit

@njit
def _numba1(sidx, c1D):
    out = []
    n = len(sidx)
    start = 0
    grpID = []
    for i in range(1,n):
        if c1D[sidx[i]]!=c1D[sidx[i-1]]:
            out.append(sidx[start:i])
            grpID.append(c1D[sidx[start]])
            start = i
    out.append(sidx[start:])
    grpID.append(c1D[sidx[start]])
    return grpID,out

def numba1(cubes):
    c1D = np.ravel_multi_index(cubes.T, cubes.max(0)+1)
    sidx = c1D.argsort()
    out = dict(zip(*_numba1(sidx, c1D)))
    return out

Et deuxième numbasolution comme:

from numba import types
from numba.typed import Dict

int_array = types.int64[:]

@njit
def _numba2(sidx, c1D):
    n = len(sidx)
    start = 0
    outt = Dict.empty(
        key_type=types.int64,
        value_type=int_array,
    )
    for i in range(1,n):
        if c1D[sidx[i]]!=c1D[sidx[i-1]]:
            outt[c1D[sidx[start]]] = sidx[start:i]
            start = i
    outt[c1D[sidx[start]]] = sidx[start:]
    return outt

def numba2(cubes):
    c1D = np.ravel_multi_index(cubes.T, cubes.max(0)+1)    
    sidx = c1D.argsort()
    out = _numba2(sidx, c1D)
    return out

Horaires avec cubes.npzdonnées -

In [4]: cubes = np.load('cubes.npz')['array']

In [5]: %timeit numpy1(cubes)
   ...: %timeit numba1(cubes)
   ...: %timeit numba2(cubes)
2.38 s ± 14.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
2.13 s ± 25.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.8 s ± 5.95 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Alternative n ° 1: nous pouvons accélérer davantage le numexprcalcul pour les grands tableaux c1D, comme ceci -

import numexpr as ne

s0,s1 = cubes[:,0].max()+1,cubes[:,1].max()+1
d = {'s0':s0,'s1':s1,'c0':cubes[:,0],'c1':cubes[:,1],'c2':cubes[:,2]}
c1D = ne.evaluate('c0+c1*s0+c2*s0*s1',d)

Ce serait applicable à tous les endroits qui le nécessitent c1D.

Vous pouvez simplement parcourir et ajouter l'index de chaque élément à la liste correspondante.

from collections import defaultdict

res = defaultdict(list)

for idx, elem in enumerate(cubes):
    #res[tuple(elem)].append(idx)
    res[elem.tobytes()].append(idx)

L'exécution peut être encore améliorée en utilisant tobytes () au lieu de convertir la clé en tuple.

Vous pouvez utiliser Cython:

%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True

import math
import cython as cy

cimport numpy as cnp


cpdef groupby_index_dict_cy(cnp.int32_t[:, :] arr):
    cdef cy.size_t size = len(arr)
    result = {}
    for i in range(size):
        key = arr[i, 0], arr[i, 1], arr[i, 2]
        if key in result:
            result[key].append(i)
        else:
            result[key] = [i]
    return result

mais cela ne vous rendra pas plus rapide que ce que Pandas fait, bien qu'il soit le plus rapide après cela (et peut-être la numpy_indexsolution basée), et ne vient pas avec la pénalité de mémoire de celui-ci. Une collection de ce qui a été proposé jusqu'à présent est ici .

Dans la machine OP, le temps d'exécution devrait atteindre près de 12 secondes.