OctreeManager and vector operations

octreelib/grid/grid.py

@@ -1,5 +1,6 @@
+from dataclasses import dataclass, field
 import random
-from typing import List, Dict, Callable
+from typing import List, Dict, Callable, Optional, Type
 
 import k3d
 import numpy as np
@@ -10,19 +11,30 @@
     GridVisualizationType,
     VisualizationConfig,
 )
-from octreelib.internal.point import (
-    RawPointCloud,
-    RawPoint,
-    PointCloud,
-)
-from octreelib.internal.voxel import Voxel
-from octreelib.octree.multi_pose_octree import MultiPoseOctree
+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"]
 
 
+@dataclass
 class GridConfig(GridConfigBase):
-    _compatible_octree_types = [MultiPoseOctree]
+    """
+    Config for Grid
+
+    octree_manager_type: type of OctreeManager used.
+        OctreeManager is responsible for managing octrees for different poses.
+    octree_type: type of Octree used
+        Octree is the data structure used for storing points.
+    octree_config: This config will be forwarded to the octrees.
+    debug: True enables debug mode.
+    voxel_edge_length: Initial size of voxels.
+    corner: Corner of a grid.
+    """
+
+    pass
 
 
 class Grid(GridBase):
@@ -38,105 +50,120 @@ class Grid(GridBase):
     def __init__(self, grid_config: GridConfig):
         super().__init__(grid_config)
 
-        # {pose -> list of voxel coordinates}
-        self.__pose_voxel_coordinates: Dict[int, List[RawPoint]] = {}
+        # {pose -> list of voxels}
+        self.__pose_voxel_coordinates: Dict[int, List[VoxelBase]] = {}
 
-        # {voxel coordinates hash -> octree}
-        self.__octrees: Dict[int, grid_config.octree_type] = {}
+        # {voxel -> octree manager}
+        self.__octrees: Dict[VoxelBase, grid_config.octree_manager_type] = {}
 
-    def insert_points(self, pose_number: int, points: RawPointCloud):
+    def insert_points(self, pose_number: int, points: PointCloud):
         """
-        Insert points to the grid
-        :param pose_number: pose to which the cloud is inserted
-        :param points: point cloud
+        Insert points to the according octree.
+        If an octree for this pos 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
         self.__pose_voxel_coordinates[pose_number] = []
 
-        # convert points to PointsWithPose, which is a subclass of np.ndarray
-        points = PointCloud(points)
-
-        for point in points:
-            # get coords of voxel into which the point is inserted
-            voxel_coordinates = self.__get_voxel_for_point(point)
-            voxel_coordinates_hash = hash(voxel_coordinates)
-
-            # create octree in the voxel if it does not exist yet
-            if voxel_coordinates_hash not in self.__octrees:
-                self.__octrees[voxel_coordinates_hash] = self._grid_config.octree_type(
+        # distribute points to voxels
+        voxel_indices = (
+            ((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():
+            target_voxel = VoxelBase(
+                np.array(voxel_coordinates),
+                self._grid_config.voxel_edge_length,
+            )
+            if target_voxel not in self.__octrees:
+                self.__octrees[target_voxel] = self._grid_config.octree_manager_type(
+                    self._grid_config.octree_type,
                     self._grid_config.octree_config,
-                    voxel_coordinates,
-                    self._grid_config.grid_voxel_edge_length,
+                    np.array(voxel_coordinates),
+                    self._grid_config.voxel_edge_length,
                 )
 
-            self.__octrees[voxel_coordinates_hash].insert_points(
-                [point.with_pose(pose_number)]
-            )
+            self.__pose_voxel_coordinates[pose_number].append(target_voxel)
+            self.__octrees[target_voxel].insert_points(pose_number, voxel_points)
 
-    def map_leaf_points(self, function: Callable[[RawPointCloud], RawPointCloud]):
+    def map_leaf_points(self, function: Callable[[PointCloud], PointCloud]):
         """
-        Transforms point cloud in each leaf node of each octree using the function
-        :param function: transformation function PointCloud -> PointCloud
+        Transform point cloud in each node of each octree using the function
+        :param function: Transformation function PointCloud -> PointCloud. It is applied to each leaf node.
         """
         for voxel_coordinates in self.__octrees:
             self.__octrees[voxel_coordinates].map_leaf_points(function)
 
     def get_leaf_points(self, pose_number: int) -> List[Voxel]:
         """
-        :param pose_number: Pose number.
+        :param pose_number: The desired pose number.
         :return: List of voxels. Each voxel is a representation of a leaf node.
         Each voxel has the same corner, edge_length and points as one of the leaf nodes.
         """
         return sum(
             [
-                octree.get_leaf_points_by_pose_number(pose_number)
-                for octree in self.__octrees.values()
+                self.__octrees[voxel_coordinates].get_leaf_points(pose_number)
+                for voxel_coordinates in self.__pose_voxel_coordinates[pose_number]
             ],
             [],
         )
 
-    def get_points(self, pose_number: int) -> RawPointCloud:
+    def get_points(self, pose_number: int) -> PointCloud:
         """
-        :param pose_number: Pose number.
-        :return: All points inside the grid.
+        Returns points for a specific pose number.
+        :param pose_number: The desired pose number.
+        :return: Points belonging to the pose.
         """
         return np.vstack(
-            [
-                octree.get_points_by_pose_number(pose_number)
-                for octree in self.__octrees.values()
-            ]
+            [octree.get_points(pose_number) for octree in self.__octrees.values()]
         )
 
-    def subdivide(self, subdivision_criteria: List[Callable[[RawPointCloud], bool]]):
+    def subdivide(
+        self,
+        subdivision_criteria: List[Callable[[PointCloud], bool]],
+        pose_numbers: Optional[List[int]] = None,
+    ):
         """
         Subdivides all octrees based on all points and given subdivision criteria.
+        :param pose_numbers: List of pose numbers to subdivide.
         :param subdivision_criteria: List of bool functions which represent criteria for subdivision.
         If any of the criteria returns **true**, the octree node is subdivided.
         """
         for voxel_coordinates in self.__octrees:
-            self.__octrees[voxel_coordinates].subdivide(subdivision_criteria)
+            self.__octrees[voxel_coordinates].subdivide(
+                subdivision_criteria, pose_numbers
+            )
 
-    def filter(self, filtering_criteria: List[Callable[[RawPointCloud], bool]]):
+    def filter(self, filtering_criteria: List[Callable[[PointCloud], bool]]):
         """
-        Filters nodes of each octree with points by criteria
-        :param filtering_criteria: Filtering Criteria
+        Filters nodes of each octree with points by criterion
+        :param filtering_criteria: List of bool functions which represent criteria for filtering.
+            If any of the criteria returns **false**, the point cloud in octree leaf is removed.
         """
         for voxel_coordinates in self.__octrees:
             self.__octrees[voxel_coordinates].filter(filtering_criteria)
 
-    def visualize(self, config: VisualizationConfig) -> None:
+    def visualize(self, config: VisualizationConfig = VisualizationConfig()) -> None:
         """
         Produces `.html` file with Grid
         """
         plot = k3d.Plot()
         random.seed(config.seed)
-        poses_number = len(self.__octrees.keys()) + 1
+        poses_numbers = self.__pose_voxel_coordinates.keys()
 
         if config.type is GridVisualizationType.POSE:
-            for pose_number in range(poses_number):
+            for pose_number in poses_numbers:
                 color = random.randrange(0, 0xFFFFFF)
                 points = self.get_points(pose_number=pose_number)
 
@@ -147,7 +174,7 @@ def visualize(self, config: VisualizationConfig) -> None:
                 )
         elif config.type is GridVisualizationType.VOXEL:
             voxels_colors = {}
-            for pose_number in range(poses_number):
+            for pose_number in poses_numbers:
                 leaves = self.get_leaf_points(pose_number=pose_number)
                 for leaf in leaves:
                     if leaf.id not in voxels_colors.keys():
@@ -161,7 +188,7 @@ def visualize(self, config: VisualizationConfig) -> None:
                     )
 
         vertices = []
-        for pose_number in range(poses_number):
+        for pose_number in poses_numbers:
             # TODO: Draw full grid including empty voxels
             leaves = self.get_leaf_points(pose_number=pose_number)
             for leaf in leaves:
@@ -190,47 +217,21 @@ def visualize(self, config: VisualizationConfig) -> None:
 
     def n_leaves(self, pose_number: int) -> int:
         """
-        :param pose_number: Pose number.
-        :return: Number of leaves in all octrees, which store points for given pose.
+        :param pose_number: The desired pose number.
+        :return: Number of leaf nodes in the octree for given pose number.
         """
-        return sum(
-            [
-                octree.n_leaves_by_pose_number(pose_number)
-                for octree in self.__octrees.values()
-            ]
-        )
+        return sum([octree.n_leaves(pose_number) for octree in self.__octrees.values()])
 
     def n_points(self, pose_number: int) -> int:
         """
-        :param pose_number: Pose number.
-        :return: Number of points for given pose.
+        :param pose_number: The desired pose number.
+        :return: Number of points of an octree for given pose number.
         """
-        return sum(
-            [
-                octree.n_points_by_pose_number(pose_number)
-                for octree in self.__octrees.values()
-            ]
-        )
+        return sum([octree.n_points(pose_number) for octree in self.__octrees.values()])
 
     def n_nodes(self, pose_number: int) -> int:
         """
-        :param pose_number: Pose number.
-        :return: Number of nodes in all octrees, which store points for given pose
-        (either themselves, or through their child nodes).
-        """
-        return sum(
-            [
-                octree.n_nodes_by_pose_number(pose_number)
-                for octree in self.__octrees.values()
-            ]
-        )
-
-    def __get_voxel_for_point(self, point: RawPoint) -> RawPoint:
-        """
-        Method to get coordinates of a voxel where the given point would be stored.
-        :param point: Point.
-        :return: Corner of the voxel in the grid, where an appropriate octree for the point resides.
+        :param pose_number: The desired pose number.
+        :return: Number of nodes of an octree for given pose number.
         """
-        point = point[:3]
-        grid_voxel_edge_length = self._grid_config.grid_voxel_edge_length
-        return point // grid_voxel_edge_length * grid_voxel_edge_length
+        return sum([octree.n_nodes(pose_number) for octree in self.__octrees.values()])