在numba中存储和处理点云数据的有效方法

Question

我需要存储一些类似这样的数据结构：

{'x1,y1,z1': [[p11_x,p11_y,p11_z], [p12_x,p12_y,p12_z], ..., [p1n_x,p1n_y,p1n_z]],
 'x2,y2,z2': [[p21_x,p21_y,p21_z], [p22_x,p22_y,p22_z], ..., [p2n_x,p2n_y,p2n_z]],
 ...
 'xn,yn,zn': [[pn1_x,pn1_y,pn1_z], [pn2_x,pn2_y,pn2_z], ..., [pnm_x,pnm_y,pnm_z]]}

每个键是网格单元索引，值是分类点的列表。列表可以是可变长度的，但我可以将它设置为静态的，例如。1000个元素。

现在，我尝试了下面这样的方法：

np.zeros(shape=(100,100,100,50,3))

但是如果我在这个函数中使用numba.jit，它产生的最坏时间是纯python的几倍。我想做的简单的python示例：

def split_into_grid_py(points: np.array):
    grid = {}
    for point in points:
        r_x = round(point[0])
        r_y = round(point[1])
        r_z = round(point[2])
        try:
            grid[(r_x, r_y, r_z)].append(point)
        except KeyError:
            grid[(r_x, r_y, r_z)] = [point]
    return grid

使用numba有什么有效的方法可以做到这一点吗？每10个循环的执行时间是这样的: numba 7.050494909286499纯python 1.0014197826385498与相同的数据集，所以这是垃圾优化。

还有我的numba代码：

@numba.jit(nopython=True)
def split_into_grid(points: np.array):
    grid = np.zeros(shape=(100,100,100,50,3))
    for point in points:
        r_x = round(point[0])
        r_y = round(point[1])
        r_z = round(point[2])
        i = 0
        for cell in grid[r_x][r_y][r_z]:
            if not np.sum(cell):
               grid[r_x][r_y][r_z][i] = point
               break
            i += 1
    return grid

Answer 1

纯Python版本使用字典容器在O(1)时间内追加条目，而Numba版本使用O(n)数组搜索(以50为界)。此外，np.zeros(shape=(100,100,100,50,3))分配了大约1个GiB的数组，这导致在将在内存中完成的计算过程中许多缓存未命中。同时，纯Python版本可能适合CPU缓存。有两种策略可以解决这个问题。

第一种策略是使用3个容器。将每个网格单元映射到第二个数组valueGrid中的偏移量的数组keyGrid，如果没有与该单元关联的点，则为-1。valueGrid包含给定网格单元的所有点。最后，countingGrid计算每个网格单元的点数。下面是一个未经测试的示例：

@numba.jit(nopython=True)
def split_into_grid(points: np.array):
    # Note: use np.uint16 if the actual number of filled grid cell is less than 65536
    keyGrid = np.full(shape=(100,100,100), -1, dtype=np.uint32)
    i = 0
    for point in points:
        r_x = round(point[0])
        r_y = round(point[1])
        r_z = round(point[2])
        if keyGrid[r_x,r_y,r_z] < 0:
            keyGrid[r_x,r_y,r_z] = i
            i += 1
    uniqueCloundPointCount = i
    # Note the number of points per grid cell is also bounded by the type
    countingGrid = np.zeros(uniqueCloundPointCount, dtype=np.uint8)
    valueGrid = np.full((uniqueCloundPointCount, 50, 3), -1, dtype=np.int32)
    for point in points:
        r_x = round(point[0])
        r_y = round(point[1])
        r_z = round(point[2])
        key = keyGrid[r_x,r_y,r_z]
        addingPos = countingGrid[key]
        valueGrid[key, addingPos] = point
        countingGrid[key] += 1
    return (keyGrid, valueGrid, countingGrid)

请注意，只要不是所有网格单元都包含点，这些数组就非常小，从而导致较少的缓存未命中。此外，每个点的映射都是在(很小的)恒定时间内完成的，因此代码速度要快得多。

第二种策略是使用与纯Python实现相同的方法，但使用Numba类型。事实上，Numba实验性地supports dictionaries。您可以将异常替换为字典检查((r_x, r_y, r_z) in grid)，这将导致较少的编译问题，并可能加速生成的代码。注意，Numba dict通常和CPython一样快(如果不是更慢的话)。因此，生成的代码可能不会快很多。