什么是将numpy数组的组名映射到索引的最快方法？

Question

我正在处理激光雷达的3D点云。这些点是通过numpy数组给出的，看起来像这样：

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]])

我希望将我的数据分组为大小为50*50*50的立方体，以便每个立方体都保留一些可哈希索引和包含在其中的points的numpy索引。为了进行分割，我使用cubes = points \\ 50进行赋值，输出结果如下：

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]])

我期望的输出如下：

{(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]}

我的真实点云包含几亿个三维点。如何以最快的方式进行这种分组？

我尝试了大多数不同的解决方案。以下是时间消耗的比较，假设点的数量约为2000万，不同立方体的大小约为100万：

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

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

defaultdict [elem.tobytes() or 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]

：将整数数组转换为NumPy数组。

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

使用Pandas和降维技术[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

你可以在这里下载cubes.npz文件，并使用命令

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

检查性能时间。

- mathfux

---

这样一个简单的Pandas解决方案很可能不会被简单的方法所超越，因为已经花费了大量精力进行优化。基于Cython的方法可能可以接近它，但我怀疑它是否能够胜过它。 - norok2

@mathfux 你需要将最终输出作为字典吗？还是将组和它们的索引作为两个输出也可以？ - Divakar

@Divakar 不错的观点！我更喜欢将输出作为一个字典，其中键是由3个坐标组成的元组。但这并非必要。让它成为任何保留第二个输出为值的字典即可。顺便说一句，我保持一个命令 dict(enumerate(values)) 是无成本的，因为它只需要在我的笔记本电脑上花费0.15秒。 - mathfux

@mathfux，另外我不确定你的“values”是什么。在预期的字典输出中，我看到键为三维坐标，索引为值。你能说明一下你之前说的“不是必要的”具体指什么吗？是否有其他输出格式 - 字典或非字典格式可以适合你？从时间上看，将其转换为字典是瓶颈所在。因此，这方面的澄清可能会对解决方案产生很大的影响。 - Divakar

我的意思是，键不一定是由3个坐标组成的元组。它可以是任何可哈希的东西。 - mathfux

Answer 1

每个组固定数量的索引

方法 #1

我们可以进行降维处理，将 立方体 缩减为一个1D数组。这是基于将给定的立方体数据映射到一个n维网格上来计算线性索引等价物的方法，详见 这里。然后，基于这些线性索引的唯一性，我们可以分离出唯一的组和它们对应的索引。因此，按照这些策略，我们会有一个解决方案，如下所示 -

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))

方案一: 如果cubes中的整数值太大，我们可能希望进行降维操作，选择较短的维度作为主轴。因此，在这些情况下，我们可以修改降维步骤以获取c1D，如下所示 -

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

方法二

接下来，我们可以使用Cython 动力的 kd 树进行快速最近邻查找，以获取最近邻索引，从而解决我们的问题 -

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]

通用案例：每组具有可变数量的索引

我们将通过一些拆分来扩展基于argsort的方法，以获得所需的输出，如下所示-

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))

使用一维的魔方组合作为键

我们将通过使用魔方组合作为键来扩展早期列出的方法，从而简化字典创建的过程，并使其更加高效，如下所示 -

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

接下来，我们将利用numba包来迭代并获得最终的可哈希字典输出。有两种解决方案-一种使用numba分别获取键和值，主调用将进行zip并转换为字典；另一种则会创建一个numba-supported字典类型，因此主调用函数不需要额外的工作。

因此，我们首先有一个numba的解决方案:

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

第二种 numba 解决方案如下：

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

使用cubes.npz数据进行的时间测量 -

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)

备选方案 #1：我们可以通过 numexpr 在大数组上实现更快的速度来计算c1D，如下所示 -

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)

这将适用于所有需要 c1D 的地方。

- Divakar

---

非常感谢您的回复！我没想到在这里可以使用cKDTree。然而，您的#Approach1仍存在一些问题。输出长度仅为915791。我猜测这是dtypes int32和int64之间的某种冲突。 - mathfux

@mathfux 我假设“每个组的索引数量将是一个常数”，这是我从评论中收集到的信息。这个假设是否安全？另外，您是否正在测试cubes.npz以获得915791的输出？ - Divakar

是的，我没有测试每个组的索引数量，因为组名称的顺序可能会有所不同。我只测试了从 cubes.npz 输出的字典长度，结果为 983234，对于我提出的其他方法也是如此。 - mathfux

@mathfux 请查看针对可变数量索引的通用情况的 方法3。 - Divakar

@mathfux 是的，如果最小值小于0，通常需要进行偏移。精度方面注意得很好！ - Divakar

Answer 2

你可以迭代并将每个元素的索引添加到相应的列表中。

from collections import defaultdict

res = defaultdict(list)

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

通过使用tobytes()而不是将密钥转换为元组，可以进一步提高运行时性能。

- abc

---

我目前正在尝试对性能时间进行评估（针对2000万个数据点）。看起来我的解决方案在时间上更有效率，因为避免了迭代。但是我同意，内存消耗是巨大的。 - mathfux

另一个提议 res[tuple(elem)].append(idx) 耗时50秒，而其修改版 res[elem[0], elem[1], elem[2]].append(idx) 只需30秒。 - mathfux

Answer 3

你可以使用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

但是它不会比Pandas更快，尽管它在Pandas之后（以及也许是基于numpy_index的解决方案）是最快的，并且不带有内存惩罚。 目前提出的集合在这里。 在OP的机器上执行时间应该接近12秒。

- norok2

---

非常感谢，我稍后会进行测试。 - mathfux