Merge pull request #19 from prime-slam/point-distribution-optimizaton

fix grid, octree: faster point distribution

octreelib/grid/grid.py

@@ -1,6 +1,6 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 import random
-from typing import List, Dict, Callable, Optional, Type
+from typing import List, Dict, Callable, Optional
 
 import k3d
 import numpy as np
@@ -11,10 +11,8 @@
     GridVisualizationType,
     VisualizationConfig,
 )
-from octreelib.internal.typing import T
 from octreelib.internal.point import PointCloud
 from octreelib.internal.voxel import Voxel, VoxelBase
-from octreelib.octree import Octree, OctreeConfig
 
 __all__ = ["Grid", "GridConfig"]
 
@@ -59,29 +57,41 @@ def __init__(self, grid_config: GridConfig):
     def insert_points(self, pose_number: int, points: PointCloud):
         """
         Insert points to the according octree.
-        If an octree for this pos does not exist, a new octree is created
+        If an octree for this pose does not exist, a new octree is created
         :param pose_number: Pose number to which the cloud is inserted.
         :param points: Point cloud to be inserted.
         """
         if pose_number in self.__pose_voxel_coordinates:
             raise ValueError(f"Cannot insert points to existing pose {pose_number}")
 
-        # register pose
+        # Register pose
         self.__pose_voxel_coordinates[pose_number] = []
 
-        # distribute points to voxels
+        # Distribute points to voxels
         voxel_indices = (
-            ((points - self._grid_config.corner) // self._grid_config.voxel_edge_length)
+            (points - self._grid_config.corner)
+            // self._grid_config.voxel_edge_length
             * self._grid_config.voxel_edge_length
         ).astype(int)
-        distributed_points = {}
-        unique_indices = np.unique(voxel_indices, axis=0)
-        for unique_id in unique_indices:
-            mask = np.where((voxel_indices == unique_id).all(axis=1))
-            distributed_points[tuple(unique_id)] = points[mask]
-
-        # insert points to octrees
-        for voxel_coordinates, voxel_points in distributed_points.items():
+
+        # Create a unique identifier for each voxel based on its indices
+        unique_voxel_indices, point_inverse_indices = np.unique(
+            voxel_indices, axis=0, return_inverse=True
+        )
+
+        # Points are reordered based on the `point_inverse_indices`, so that they can be split
+        # into groups of points, where each group is inserted into the corresponding voxel.
+        # The indices for splitting are calculated using `np.cumsum()` based on the number
+        # of points which would be distributed into each voxel.
+        grouped_points = np.split(
+            points[point_inverse_indices.argsort()],
+            np.cumsum(np.bincount(point_inverse_indices))[:-1],
+        )
+
+        # Insert points to octrees
+        for voxel_coordinates, voxel_points in zip(
+            unique_voxel_indices, grouped_points
+        ):
             target_voxel = VoxelBase(
                 np.array(voxel_coordinates),
                 self._grid_config.voxel_edge_length,

octreelib/octree/octree.py

@@ -67,21 +67,33 @@ def insert_points(self, points: PointCloud):
         :param points: Points to insert.
         """
         if self._has_children:
-            # distribute points to children
-            clouds = []
-            for offset in itertools.product([0, self.edge_length / 2], repeat=3):
-                child_corner_min = self.corner_min + np.array(offset)
-                child_corner_max = child_corner_min + self.edge_length / 2
-                # find points in the child
-                mask = np.all(
-                    (points >= child_corner_min) & (points < child_corner_max), axis=1
-                )
-                child_points = points[mask]
-                clouds.append(child_points)
+            # For all points calculate the voxel to insert in
+            voxel_indices = (
+                (points - self.corner_min) // (self.edge_length / 2)
+            ).astype(int)
+
+            # Create a unique identifier for each voxel based on its indices
+            unique_voxel_indices, point_inverse_indices = np.unique(
+                voxel_indices, axis=0, return_inverse=True
+            )
 
-            # insert points to children
-            for child, cloud in zip(self._children, clouds):
-                child.insert_points(cloud)
+            # Points are reordered based on the `point_inverse_indices`, so that they can be split
+            # into groups of points, where each group is inserted into the corresponding voxel.
+            # The indices for splitting are calculated using `np.cumsum()` based on the number
+            # of points which would be distributed into each voxel.
+            grouped_points = np.split(
+                points[point_inverse_indices.argsort()],
+                np.cumsum(np.bincount(point_inverse_indices))[:-1],
+            )
+            for unique_voxel_index, child_points in zip(
+                unique_voxel_indices, grouped_points
+            ):
+                # Calculate the internal child id from its binary representation
+                child_id = sum(
+                    2**i * exists_offset
+                    for i, exists_offset in enumerate(unique_voxel_index[::-1])
+                )
+                self._children[child_id].insert_points(child_points)
         else:
             self._points = np.vstack([self._points, points])