Docstring updates for NearestNeighborSearch and compute_metrics

cpp/pybind/core/nns/nearest_neighbor_search.cpp

@@ -21,10 +21,35 @@ namespace nns {
 void pybind_core_nns_declarations(py::module &m_nns) {
     py::class_<NearestNeighborSearch, std::shared_ptr<NearestNeighborSearch>>
             nns(m_nns, "NearestNeighborSearch",
-                "NearestNeighborSearch class for nearest neighbor search. "
-                "Construct a NearestNeighborSearch object with input "
-                "dataset_points of shape {n_dataset, d}.");
+                R"(NearestNeighborSearch class for nearest neighbor search. 
+
+This class holds multiple index types to accelerate various nearest neighbor 
+search operations for a dataset of points with shape {n,d} with `n` as the number
+of points and `d` as the dimension of the points. The class supports knn search,
+fixed-radius search, multi-radius search, and hybrid search.
+
+Example:
+    The following example demonstrates how to perform knn search using the class::
+
+        import open3d as o3d
+        import numpy as np
+
+        dataset = np.random.rand(10,3)
+        query_points = np.random.rand(5,3)
+
+        nns = o3d.core.nns.NearestNeighborSearch(dataset)
+
+        # initialize the knn_index before we can use knn_search
+        nns.knn_index()
+
+        # perform knn search to get the 3 closest points with respect to the 
+        # Euclidean distance. The returned distance is given as the squared
+        # distances
+        indices, squared_distances = nns.knn_search(query_points, knn=3)
+
+                )");
 }
+
 void pybind_core_nns_definitions(py::module &m_nns) {
     auto nns = static_cast<py::class_<NearestNeighborSearch,
                                       std::shared_ptr<NearestNeighborSearch>>>(
@@ -35,7 +60,14 @@ void pybind_core_nns_definitions(py::module &m_nns) {
 
     // Index functions.
     nns.def("knn_index", &NearestNeighborSearch::KnnIndex,
-            "Set index for knn search.");
+            R"(Initialize the index for knn search. 
+
+This function needs to be called once before performing search operations.
+
+Returns:
+    True on success.
+            )");
+
     nns.def(
             "fixed_radius_index",
             [](NearestNeighborSearch &self, utility::optional<double> radius) {
@@ -45,9 +77,28 @@ void pybind_core_nns_definitions(py::module &m_nns) {
                     return self.FixedRadiusIndex(radius.value());
                 }
             },
-            py::arg("radius") = py::none());
+            py::arg("radius") = py::none(),
+            R"(Initialize the index for fixed-radius search.
+
+This function needs to be called once before performing search operations.
+
+Args:
+    radius (float, optional): Radius value for fixed-radius search. Required
+        for GPU fixed radius index.
+
+Returns:
+    True on success.
+            )");
+
     nns.def("multi_radius_index", &NearestNeighborSearch::MultiRadiusIndex,
-            "Set index for multi-radius search.");
+            R"(Initialize the index for multi-radius search.
+
+This function needs to be called once before performing search operations.
+
+Returns:
+        True on success.
+             )");
+
     nns.def(
             "hybrid_index",
             [](NearestNeighborSearch &self, utility::optional<double> radius) {
@@ -57,11 +108,57 @@ void pybind_core_nns_definitions(py::module &m_nns) {
                     return self.HybridIndex(radius.value());
                 }
             },
-            py::arg("radius") = py::none());
+            py::arg("radius") = py::none(),
+            R"(Initialize the index for hybrid search.
+
+This function needs to be called once before performing search operations.
+
+Args:
+    radius (float, optional): Radius value for hybrid search. Required
+        for GPU hybrid index.
+
+Returns:
+    True on success.
+            )");
 
     // Search functions.
     nns.def("knn_search", &NearestNeighborSearch::KnnSearch, "query_points"_a,
-            "knn"_a, "Perform knn search.");
+            "knn"_a,
+            R"(Perform knn search.
+
+Note:
+    To use knn_search initialize the index using knn_index before calling this function.
+
+Args:
+    query_points (open3d.core.Tensor): Query points with shape {n, d}.
+    knn (int): Number of neighbors to search per query point.
+
+Example:
+    The following searches the 3 nearest neighbors for random dataset and query points::
+
+        import open3d as o3d
+        import numpy as np
+
+        dataset = np.random.rand(10,3)
+        query_points = np.random.rand(5,3)
+
+        nns = o3d.core.nns.NearestNeighborSearch(dataset)
+
+        # initialize the knn_index before we can use knn_search
+        nns.knn_index()
+
+        # perform knn search to get the 3 closest points with respect to the 
+        # Euclidean distance. The returned distance is given as the squared
+        # distances
+        indices, squared_distances = nns.knn_search(query_points, knn=3)
+
+Returns:
+    Tuple of Tensors (indices, squared_distances).
+        - indices: Tensor of shape {n, knn}.
+        - squared_distances: Tensor of shape {n, knn}. The distances are squared L2 distances.
+
+            )");
+
     nns.def(
             "fixed_radius_search",
             [](NearestNeighborSearch &self, Tensor query_points, double radius,
@@ -74,38 +171,141 @@ void pybind_core_nns_definitions(py::module &m_nns) {
                 }
             },
             py::arg("query_points"), py::arg("radius"),
-            py::arg("sort") = py::none());
+            py::arg("sort") = py::none(),
+            R"(Perform fixed-radius search.
+
+Note:
+    To use fixed_radius_search initialize the index using fixed_radius_index before calling this function.
+
+Args:
+        query_points (open3d.core.Tensor): Query points with shape {n, d}.
+        radius (float): Radius value for fixed-radius search. Note that this
+            parameter can differ from the radius used to initialize the index
+            for convenience, which may cause the index to be rebuilt for GPU 
+            devices.
+        sort (bool, optional): Sort the results by distance. Default is True.
+
+Returns:
+        Tuple of Tensors (indices, splits, distances).
+            - indices: The indices of the neighbors.
+            - distances: The squared L2 distances.
+            - splits: The splits of the indices and distances defining the start 
+                and exclusive end of each query point's neighbors. The shape is {num_queries+1}
+
+Example:
+    The following searches the neighbors within a radius of 1.0 a set of data and query points::
+
+        # define data and query points
+        points = np.array([
+        [0.1,0.1,0.1],
+        [0.9,0.9,0.9],
+        [0.5,0.5,0.5],
+        [1.7,1.7,1.7],
+        [1.8,1.8,1.8],
+        [0.3,2.4,1.4]], dtype=np.float32)
+
+        queries = np.array([
+        [1.0,1.0,1.0],
+        [0.5,2.0,2.0],
+        [0.5,2.1,2.1],
+        [100,100,100],
+        ], dtype=np.float32)
+
+        nns = o3d.core.nns.NearestNeighborSearch(points)
+        nns.fixed_radius_index(radius=1.0)
+
+        neighbors_index, neighbors_distance, neighbors_splits = nns.fixed_radius_search(queries, radius=1.0, sort=True)
+        # returns neighbors_index      = [1, 2, 5, 5]
+        #         neighbors_distance   = [0.03 0.75 0.56000006 0.62]
+        #         neighbors_splits = [0, 2, 3, 4, 4]
+
+        for i, (start_i, end_i) in enumerate(zip(neighbors_splits, neighbors_splits[1:])):
+        start_i = start_i.item()
+        end_i = end_i.item()
+        print(f"query_point {i} has the neighbors {neighbors_index[start_i:end_i].numpy()} "
+              f"with squared distances {neighbors_distance[start_i:end_i].numpy()}")
+            )");
+
     nns.def("multi_radius_search", &NearestNeighborSearch::MultiRadiusSearch,
             "query_points"_a, "radii"_a,
-            "Perform multi-radius search. Each query point has an independent "
-            "radius.");
+            R"(Perform multi-radius search. Each query point has an independent radius.
+
+Note:
+    To use multi_radius_search initialize the index using multi_radius_index before calling this function.
+
+Args:
+    query_points (open3d.core.Tensor): Query points with shape {n, d}.
+    radii (open3d.core.Tensor): Radii of query points. Each query point has one radius.
+
+Returns:
+        Tuple of Tensors (indices, splits, distances).
+            - indices: The indices of the neighbors.
+            - distances: The squared L2 distances.
+            - splits: The splits of the indices and distances defining the start 
+                and exclusive end of each query point's neighbors. The shape is {num_queries+1}
+
+Example:
+    The following searches the neighbors with an individual radius for each query point::
+
+        # define data, query points and radii
+        points = np.array([
+        [0.1,0.1,0.1],
+        [0.9,0.9,0.9],
+        [0.5,0.5,0.5],
+        [1.7,1.7,1.7],
+        [1.8,1.8,1.8],
+        [0.3,2.4,1.4]], dtype=np.float32)
+
+        queries = np.array([
+            [1.0,1.0,1.0],
+            [0.5,2.0,2.0],
+            [0.5,2.1,2.1],
+            [100,100,100],
+        ], dtype=np.float32)
+
+        radii = np.array([0.5, 1.0, 1.5, 2], dtype=np.float32)
+
+        nns = o3d.core.nns.NearestNeighborSearch(points)
+
+        nns.multi_radius_index()
+
+        neighbors_index, neighbors_distance, neighbors_splits = nns.multi_radius_search(queries, radii)
+        # returns neighbors_index      = [1 5 5 3 4]
+        #         neighbors_distance   = [0.03 0.56 0.62 1.76 1.87]
+        #         neighbors_splits = [0 1 2 5 5]
+
+
+        for i, (start_i, end_i) in enumerate(zip(neighbors_splits, neighbors_splits[1:])):
+            start_i = start_i.item()
+            end_i = end_i.item()
+            print(f"query_point {i} has the neighbors {neighbors_index[start_i:end_i].numpy()} "
+                f"with squared distances {neighbors_distance[start_i:end_i].numpy()}")
+            )");
+
     nns.def("hybrid_search", &NearestNeighborSearch::HybridSearch,
             "query_points"_a, "radius"_a, "max_knn"_a,
-            "Perform hybrid search.");
-
-    // Docstrings.
-    static const std::unordered_map<std::string, std::string>
-            map_nearest_neighbor_search_method_docs = {
-                    {"query_points", "The query tensor of shape {n_query, d}."},
-                    {"radii",
-                     "Tensor of shape {n_query,} containing multiple radii, "
-                     "one for each query point."},
-                    {"radius", "Radius value for radius search."},
-                    {"max_knn",
-                     "Maximum number of neighbors to search per query point."},
-                    {"knn", "Number of neighbors to search per query point."}};
-    docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
-                                    "knn_search",
-                                    map_nearest_neighbor_search_method_docs);
-    docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
-                                    "multi_radius_search",
-                                    map_nearest_neighbor_search_method_docs);
-    docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
-                                    "fixed_radius_search",
-                                    map_nearest_neighbor_search_method_docs);
-    docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
-                                    "hybrid_search",
-                                    map_nearest_neighbor_search_method_docs);
+            R"(Perform hybrid search.
+
+Hybrid search behaves similarly to fixed-radius search, but with a maximum number of neighbors to search per query point.
+
+Note:
+    To use hybrid_search initialize the index using hybrid_index before calling this function.
+
+Args:
+    query_points (open3d.core.Tensor): Query points with shape {n, d}.
+    radius (float): Radius value for hybrid search.
+    max_knn (int): Maximum number of neighbor to search per query.
+
+Returns:
+        Tuple of Tensors (indices, distances, counts)
+            - indices: The indices of the neighbors with shape {n, max_knn}.
+                If there are less than max_knn neighbors within the radius then the
+                last entries are padded with -1.
+            - distances: The squared L2 distances with shape {n, max_knn}.
+                If there are less than max_knn neighbors within the radius then the
+                last entries are padded with 0.
+            - counts: Counts of neighbour for each query points with shape {n}.
+            )");
 }
 
 }  // namespace nns

cpp/pybind/t/geometry/pointcloud.cpp

@@ -769,17 +769,29 @@ the partition id for each point.
 
 )");
 
-    pointcloud.def(
-            "compute_metrics", &PointCloud::ComputeMetrics, "pcd2"_a,
-            "metrics"_a, "params"_a,
-            R"(Compute various metrics between two point clouds. Currently, Chamfer distance, Hausdorff distance and F-Score <a href="../tutorial/reference.html#Knapitsch2017">[[Knapitsch2017]]</a> are supported. The Chamfer distance is the sum of the mean distance to the nearest neighbor from the points of the first point cloud to the second point cloud. The F-Score at a fixed threshold radius is the harmonic mean of the Precision and Recall. Recall is the percentage of surface points from the first point cloud that have the second point cloud points within the threshold radius, while Precision is the percentage of points from the second point cloud that have the first point cloud points within the threhold radius.
-
-    .. math:
-        \text{Chamfer Distance: } d_{CD}(X,Y) = \frac{1}{|X|}\sum_{i \in X} || x_i - n(x_i, Y) || + \frac{1}{|Y|}\sum_{i \in Y} || y_i - n(y_i, X) ||\\
-        \text{Hausdorff distance: } d_H(X,Y) = \max \left{ \max_{i \in X} || x_i - n(x_i, Y) ||, \max_{i \in Y} || y_i - n(y_i, X) || \right}\\
-        \text{Precision: } P(X,Y|d) = \frac{100}{|X|} \sum_{i \in X} || x_i - n(x_i, Y) || < d \\
-        \text{Recall: } R(X,Y|d) = \frac{100}{|Y|} \sum_{i \in Y} || y_i - n(y_i, X) || < d \\
-        \text{F-Score: } F(X,Y|d) = \frac{2 P(X,Y|d) R(X,Y|d)}{P(X,Y|d) + R(X,Y|d)} \\
+    pointcloud.def("compute_metrics", &PointCloud::ComputeMetrics, "pcd2"_a,
+                   "metrics"_a, "params"_a,
+                   R"(Compute various metrics between two point clouds. 
+
+Currently, Chamfer distance, Hausdorff distance and F-Score `[Knapitsch2017] <../tutorial/reference.html#Knapitsch2017>`_ are supported. 
+The Chamfer distance is the sum of the mean distance to the nearest neighbor 
+from the points of the first point cloud to the second point cloud. The F-Score
+at a fixed threshold radius is the harmonic mean of the Precision and Recall. 
+Recall is the percentage of surface points from the first point cloud that have 
+the second point cloud points within the threshold radius, while Precision is 
+the percentage of points from the second point cloud that have the first point 
+cloud points within the threhold radius.
+
+.. math::
+    :nowrap:
+
+    \begin{align}
+        \text{Chamfer Distance: } d_{CD}(X,Y) &= \frac{1}{|X|}\sum_{i \in X} || x_i - n(x_i, Y) || + \frac{1}{|Y|}\sum_{i \in Y} || y_i - n(y_i, X) ||\\
+        \text{Hausdorff distance: } d_H(X,Y) &= \max \left\{ \max_{i \in X} || x_i - n(x_i, Y) ||, \max_{i \in Y} || y_i - n(y_i, X) || \right\}\\
+        \text{Precision: } P(X,Y|d) &= \frac{100}{|X|} \sum_{i \in X} || x_i - n(x_i, Y) || < d \\
+        \text{Recall: } R(X,Y|d) &= \frac{100}{|Y|} \sum_{i \in Y} || y_i - n(y_i, X) || < d \\
+        \text{F-Score: } F(X,Y|d) &= \frac{2 P(X,Y|d) R(X,Y|d)}{P(X,Y|d) + R(X,Y|d)} \\
+    \end{align}
 
 Args:
     pcd2 (t.geometry.PointCloud): Other point cloud to compare with.