conquer3d 0.2.5__tar.gz → 0.2.7__tar.gz

{conquer3d-0.2.5/conquer3d.egg-info → conquer3d-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: conquer3d
-Version: 0.2.5
+Version: 0.2.7
 Summary: Geometric Cuda Tool Box
 Author-email: Do Hoang Khoi <khoido8899@gmail.com>
 License-Expression: MIT

{conquer3d-0.2.5 → conquer3d-0.2.7}/conquer3d/_C.pyi

@@ -14,9 +14,13 @@ class BVH:
         """
         Query the BVH with bounding boxes or segments. Returns (query_ids, object_ids).
         """
-    def query_ray(self, ray_origins: torch.Tensor, ray_dirs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def query_point(self, query_points: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
-        Query the BVH with rays. Returns (query_ids, object_ids).
+        Find the closest leaf AABB to each query point. Returns (query_ids, object_ids, distances)
+        """
+    def query_ray(self, ray_origins: torch.Tensor, ray_dirs: torch.Tensor, max_capacity: int = 10000000) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Find all ray-AABB intersections. Returns (ray_ids, object_ids)
         """
     def query_self(self) -> tuple[torch.Tensor, torch.Tensor]:
         """
@@ -53,6 +57,10 @@ class MeshBVH(BVH):
         """
         Construct MeshBVH from Triangle AABBs.
         """
+    def build_winding_data(self, vertices: torch.Tensor, triangles: torch.Tensor) -> None:
+        """
+        Computes the hierarchical area-weighted constants for Fast Winding Number queries.
+        """
     def get_ray_intersection(self, ray_origins: torch.Tensor, ray_dirs: torch.Tensor, vertices: torch.Tensor, triangles: torch.Tensor, return_distance: bool = False) -> typing.Any:
         """
         Find all ray-triangle intersections. Returns (ray_ids, triangle_ids, intersect_points, [distances])
@@ -65,6 +73,10 @@ class MeshBVH(BVH):
         """
         Check if there are any self-intersecting triangle pairs
         """
+    def query_point(self, query_points: torch.Tensor, vertices: torch.Tensor, triangles: torch.Tensor, return_sdf: bool = False, return_prj_pts: bool = True, sign_mode: int = 0) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Find the closest triangle to each query point. Returns (query_ids, triangle_ids, projected_points, distances)
+        """
 class PGSBVH(BVH):
     def __init__(self, in_aabb_mins: torch.Tensor, in_aabb_maxs: torch.Tensor) -> None:
         """
@@ -89,6 +101,8 @@ class TriangleMesh:
         ...
     def compute_triangle_normals(self) -> None:
         ...
+    def compute_vertex_normals(self) -> None:
+        ...
     def compute_vertices_to_triangle_map(self) -> None:
         ...
     def get_edges_to_triangle_map(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -115,8 +129,16 @@ class TriangleMesh:
         """
     def is_vertex_manifold(self) -> bool:
         ...
+    def query_points(self, query_pts: torch.Tensor, return_sdf: bool = False, return_prj_pts: bool = True, sign_mode: int = 0, distance_mode: int = 0) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Query points against the mesh. Returns (query_ids, triangle_ids, projected_points, distances)
+        """
     def remove_triangles_by_mask(self, keep_mask: torch.Tensor) -> None:
         ...
+    def sample_points(self, num_points: int, uniform: bool = False, return_normals: bool = False, return_colors: bool = False, use_triangle_normal: bool = True) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """
+        Sample points on the mesh
+        """
     @property
     def bvh(self) -> MeshBVH:
         ...

{conquer3d-0.2.5 → conquer3d-0.2.7}/conquer3d/csrc/binds/data_structure/bvh.cpp

@@ -155,48 +155,81 @@ std::tuple<torch::Tensor, torch::Tensor> BVH::query_self()
 
 std::tuple<torch::Tensor, torch::Tensor> BVH::query_ray(
     const torch::Tensor &ray_origins,
-    const torch::Tensor &ray_dirs)
+    const torch::Tensor &ray_dirs,
+    int64_t max_capacity)
 {
     CHECK_INPUT(ray_origins);
     CHECK_INPUT(ray_dirs);
-    TORCH_CHECK(ray_origins.scalar_type() == torch::kFloat32, "ray_origins must be float32");
-    TORCH_CHECK(ray_dirs.scalar_type() == torch::kFloat32, "ray_dirs must be float32");
-    TORCH_CHECK(ray_origins.size(1) == 3, "ray_origins must have shape (M, 3)");
-    TORCH_CHECK(ray_dirs.size(1) == 3, "ray_dirs must have shape (M, 3)");
+    TORCH_CHECK(ray_origins.size(0) == ray_dirs.size(0), "ray_origins and ray_dirs must have the same number of rays");
 
-    const uint32_t num_queries = static_cast<uint32_t>(ray_origins.size(0));
-    auto options_int64 = ray_origins.options().dtype(torch::kInt64);
+    int num_queries = ray_origins.size(0);
 
-    if (num_queries == 0) {
-        throw std::runtime_error("Cannot query BVH with 0 queries.");
-    }
-
-    torch::Tensor out_query_ids = torch::empty({BVH_MAX_CAPACITY}, options_int64);
-    torch::Tensor out_object_ids = torch::empty({BVH_MAX_CAPACITY}, options_int64);
-    torch::Tensor hit_counter = torch::zeros({1}, options_int64);
+    auto options_i64 = torch::TensorOptions().dtype(torch::kInt64).device(ray_origins.device());
+    torch::Tensor out_query_ids = torch::empty({max_capacity}, options_i64);
+    torch::Tensor out_object_ids = torch::empty({max_capacity}, options_i64);
+    torch::Tensor hit_counter = torch::zeros({1}, options_i64);
 
     bvh::query_ray(
         num_queries,
         this->num_objects,
-        reinterpret_cast<const float3 *>(ray_origins.data_ptr<float>()),
-        reinterpret_cast<const float3 *>(ray_dirs.data_ptr<float>()),
-        reinterpret_cast<const float3 *>(this->aabb_mins.data_ptr<float>()),
-        reinterpret_cast<const float3 *>(this->aabb_maxs.data_ptr<float>()),
-        reinterpret_cast<const int2 *>(this->bvh_children.data_ptr<int>()),
-        reinterpret_cast<const int *>(this->object_ids.data_ptr<int>()),
-        reinterpret_cast<int64_t *>(out_query_ids.data_ptr<int64_t>()),
-        reinterpret_cast<int64_t *>(out_object_ids.data_ptr<int64_t>()),
-        reinterpret_cast<int64_t *>(hit_counter.data_ptr<int64_t>()),
-        static_cast<int64_t>(BVH_MAX_CAPACITY)
-    );
+        (const float3 *)ray_origins.data_ptr<float>(),
+        (const float3 *)ray_dirs.data_ptr<float>(),
+        (const float3 *)this->aabb_mins.data_ptr<float>(),
+        (const float3 *)this->aabb_maxs.data_ptr<float>(),
+        (const int2 *)this->bvh_children.data_ptr<int>(),
+        this->object_ids.data_ptr<int>(),
+        out_query_ids.data_ptr<int64_t>(),
+        out_object_ids.data_ptr<int64_t>(),
+        hit_counter.data_ptr<int64_t>(),
+        max_capacity);
+
+    int64_t h_hit_counter = hit_counter.item<int64_t>();
+
+    if (h_hit_counter >= max_capacity)
+    {
+        TORCH_WARN("BVH ray query capacity exceeded. Some hits were dropped.");
+        h_hit_counter = max_capacity;
+    }
 
-    int64_t num_hits = hit_counter.item<int64_t>();
-    num_hits = std::min(num_hits, static_cast<int64_t>(BVH_MAX_CAPACITY));
+    if (h_hit_counter == 0) {
+        return std::make_tuple(
+            torch::empty({0}, options_i64),
+            torch::empty({0}, options_i64)
+        );
+    }
 
     return std::make_tuple(
-        out_query_ids.slice(0, 0, num_hits), 
-        out_object_ids.slice(0, 0, num_hits)
-    );
+        out_query_ids.slice(0, 0, h_hit_counter),
+        out_object_ids.slice(0, 0, h_hit_counter));
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> BVH::query_point(
+    const torch::Tensor &query_points)
+{
+    CHECK_INPUT(query_points);
+
+    int num_queries = query_points.size(0);
+    
+    auto options_i64 = torch::TensorOptions().dtype(torch::kInt64).device(query_points.device());
+    auto options_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(query_points.device());
+    
+    torch::Tensor out_query_ids = torch::empty({num_queries}, options_i64);
+    torch::Tensor out_object_ids = torch::empty({num_queries}, options_i64);
+    torch::Tensor out_distances = torch::empty({num_queries}, options_f32);
+
+    bvh::query_point(
+        num_queries,
+        this->num_objects,
+        (const float3 *)query_points.data_ptr<float>(),
+        (const float3 *)this->aabb_mins.data_ptr<float>(),
+        (const float3 *)this->aabb_maxs.data_ptr<float>(),
+        (const int2 *)this->bvh_children.data_ptr<int>(),
+        this->object_ids.data_ptr<int>(),
+        out_query_ids.data_ptr<int64_t>(),
+        out_object_ids.data_ptr<int64_t>(),
+        out_distances.data_ptr<float>());
+
+    return std::make_tuple(out_query_ids, out_object_ids, out_distances);
 }
 
 void bind_ds_bvh(py::module_& m)
@@ -216,7 +249,11 @@ void bind_ds_bvh(py::module_& m)
              "Query the BVH against itself. Returns unique overlapping (query_ids, object_ids).")
 
         .def("query_ray", &BVH::query_ray,
-             py::arg("ray_origins"), 
+             py::arg("ray_origins"),
              py::arg("ray_dirs"),
-             "Query the BVH with rays. Returns (query_ids, object_ids).");
+             py::arg("max_capacity") = BVH_MAX_CAPACITY,
+             "Find all ray-AABB intersections. Returns (ray_ids, object_ids)")
+        .def("query_point", &BVH::query_point,
+             py::arg("query_points"),
+             "Find the closest leaf AABB to each query point. Returns (query_ids, object_ids, distances)");
 }

conquer3d-0.2.7/conquer3d/csrc/binds/data_structure/mesh_bvh.cpp

@@ -0,0 +1,238 @@
+#include <torch/extension.h>
+#include "../../data_structure/mesh_bvh.h"
+#include <pybind11/pybind11.h>
+#include "../../check.h"
+
+namespace py = pybind11;
+
+torch::Tensor MeshBVH::get_self_intersection(
+    const torch::Tensor &vertices,
+    const torch::Tensor &triangles)
+{
+     // 1. Broad-phase AABB overlap query
+     auto [query_ids, object_ids] = this->query_self();
+
+     int num_pairs = query_ids.size(0);
+     if (num_pairs == 0)
+     {
+          return torch::empty({0, 2}, torch::TensorOptions().dtype(torch::kInt64).device(query_ids.device()));
+     }
+
+     // 2. Narrow-phase intersection
+     auto options_i64 = torch::TensorOptions().dtype(torch::kInt64).device(query_ids.device());
+     torch::Tensor out_query_ids = torch::empty({num_pairs}, options_i64);
+     torch::Tensor out_object_ids = torch::empty({num_pairs}, options_i64);
+     torch::Tensor valid_counter = torch::zeros({1}, options_i64);
+
+     mesh_bvh::filter_self_intersections(
+         num_pairs,
+         query_ids.data_ptr<int64_t>(),
+         object_ids.data_ptr<int64_t>(),
+         (const float3 *)vertices.data_ptr<float>(),
+         (const int3 *)triangles.data_ptr<int>(),
+         out_query_ids.data_ptr<int64_t>(),
+         out_object_ids.data_ptr<int64_t>(),
+         valid_counter.data_ptr<int64_t>());
+
+     int64_t h_valid_counter = valid_counter.item<int64_t>();
+
+     if (h_valid_counter == 0)
+     {
+          return torch::empty({0, 2}, options_i64);
+     }
+
+     out_query_ids = out_query_ids.slice(0, 0, h_valid_counter).unsqueeze(1);
+     out_object_ids = out_object_ids.slice(0, 0, h_valid_counter).unsqueeze(1);
+
+     return torch::cat({out_query_ids, out_object_ids}, 1);
+}
+
+bool MeshBVH::is_self_intersection(
+    const torch::Tensor &vertices,
+    const torch::Tensor &triangles)
+{
+     torch::Tensor pairs = this->get_self_intersection(vertices, triangles);
+     return pairs.size(0) > 0;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> MeshBVH::get_ray_intersection(
+    const torch::Tensor &ray_origins,
+    const torch::Tensor &ray_dirs,
+    const torch::Tensor &vertices,
+    const torch::Tensor &triangles,
+    bool return_distance)
+{
+     // 1. Broad-phase BVH query
+     auto [query_ids, object_ids] = this->query_ray(ray_origins, ray_dirs);
+
+     int num_pairs = query_ids.size(0);
+     auto options_i64 = torch::TensorOptions().dtype(torch::kInt64).device(query_ids.device());
+     auto options_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(query_ids.device());
+
+     if (num_pairs == 0)
+     {
+          return std::make_tuple(
+              torch::empty({0}, options_i64),
+              torch::empty({0}, options_i64),
+              torch::empty({0, 3}, options_f32),
+              torch::empty({0}, options_f32));
+     }
+
+     // 2. Narrow-phase intersection
+     torch::Tensor out_query_ids = torch::empty({num_pairs}, options_i64);
+     torch::Tensor out_object_ids = torch::empty({num_pairs}, options_i64);
+     torch::Tensor out_intersect_pts = torch::empty({num_pairs, 3}, options_f32);
+
+     torch::Tensor out_distances;
+     if (return_distance)
+     {
+          out_distances = torch::empty({num_pairs}, options_f32);
+     }
+     else
+     {
+          out_distances = torch::empty({0}, options_f32);
+     }
+
+     torch::Tensor valid_counter = torch::zeros({1}, options_i64);
+
+     mesh_bvh::filter_ray_triangle_intersections(
+         num_pairs,
+         query_ids.data_ptr<int64_t>(),
+         object_ids.data_ptr<int64_t>(),
+         (const float3 *)ray_origins.data_ptr<float>(),
+         (const float3 *)ray_dirs.data_ptr<float>(),
+         (const float3 *)vertices.data_ptr<float>(),
+         (const int3 *)triangles.data_ptr<int>(),
+         out_query_ids.data_ptr<int64_t>(),
+         out_object_ids.data_ptr<int64_t>(),
+         (float3 *)out_intersect_pts.data_ptr<float>(),
+         return_distance ? out_distances.data_ptr<float>() : nullptr,
+         return_distance,
+         valid_counter.data_ptr<int64_t>());
+
+     int64_t h_valid_counter = valid_counter.item<int64_t>();
+
+     if (h_valid_counter == 0)
+     {
+          return std::make_tuple(
+              torch::empty({0}, options_i64),
+              torch::empty({0}, options_i64),
+              torch::empty({0, 3}, options_f32),
+              torch::empty({0}, options_f32));
+     }
+
+     return std::make_tuple(
+         out_query_ids.slice(0, 0, h_valid_counter),
+         out_object_ids.slice(0, 0, h_valid_counter),
+         out_intersect_pts.slice(0, 0, h_valid_counter),
+         return_distance ? out_distances.slice(0, 0, h_valid_counter) : out_distances);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> MeshBVH::query_point(
+    const torch::Tensor &query_points,
+    const torch::Tensor &vertices,
+    const torch::Tensor &triangles,
+    bool return_sdf,
+    bool return_prj_pts,
+    int sign_mode)
+{
+     if (sign_mode != 0 && sign_mode != 1)
+     {
+          throw std::runtime_error("sign_mode must be 0 (ray casting) or 1 (fast winding number)");
+     }
+     
+     if (sign_mode == 1 && !this->has_winding_data)
+     {
+          this->build_winding_data(vertices, triangles);
+     }
+
+     int num_queries = query_points.size(0);
+     int num_objects = this->object_ids.size(0);
+
+     auto options_i64 = torch::TensorOptions().dtype(torch::kInt64).device(query_points.device());
+     auto options_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(query_points.device());
+
+     torch::Tensor out_query_ids = torch::empty({num_queries}, options_i64);
+     torch::Tensor out_object_ids = torch::empty({num_queries}, options_i64);
+     torch::Tensor out_projected_pts;
+     if (return_prj_pts)
+     {
+          out_projected_pts = torch::empty({num_queries, 3}, options_f32);
+     }
+     else
+     {
+          out_projected_pts = torch::empty({0, 3}, options_f32);
+     }
+     torch::Tensor out_distances = torch::empty({num_queries}, options_f32);
+
+     mesh_bvh::query_point_mesh_bvh(
+         num_queries,
+         num_objects,
+         (const float3 *)query_points.data_ptr<float>(),
+         (const float3 *)vertices.data_ptr<float>(),
+         (const int3 *)triangles.data_ptr<int>(),
+         (const float3 *)this->aabb_mins.data_ptr<float>(),
+         (const float3 *)this->aabb_maxs.data_ptr<float>(),
+         (const int2 *)this->bvh_children.data_ptr<int>(),
+         this->object_ids.data_ptr<int>(),
+         this->has_winding_data ? (const WindingData *)this->winding_data.data_ptr() : nullptr,
+         out_query_ids.data_ptr<int64_t>(),
+         out_object_ids.data_ptr<int64_t>(),
+         return_prj_pts ? (float3 *)out_projected_pts.data_ptr<float>() : nullptr,
+         out_distances.data_ptr<float>(),
+         return_sdf,
+         return_prj_pts,
+         sign_mode);
+
+     return std::make_tuple(out_query_ids, out_object_ids, out_projected_pts, out_distances);
+}
+
+void MeshBVH::build_winding_data(
+    const torch::Tensor &vertices,
+    const torch::Tensor &triangles)
+{
+     int num_objects = this->object_ids.size(0);
+     int num_nodes = num_objects * 2 - 1;
+
+     auto options = torch::TensorOptions().dtype(torch::kByte).device(vertices.device());
+     this->winding_data = torch::empty({num_nodes * (int64_t)sizeof(WindingData)}, options);
+
+     mesh_bvh::bottom_up_winding_data(
+         num_objects,
+         this->object_ids.data_ptr<int>(),
+         (const float3 *)vertices.data_ptr<float>(),
+         (const int3 *)triangles.data_ptr<int>(),
+         this->bvh_parents.data_ptr<int>(),
+         (const int2 *)this->bvh_children.data_ptr<int>(),
+         (WindingData *)this->winding_data.data_ptr());
+
+     cudaDeviceSynchronize();
+     this->has_winding_data = true;
+}
+
+void bind_ds_mesh_bvh(py::module_ &m)
+{
+     py::class_<MeshBVH, BVH>(m, "MeshBVH")
+         .def(py::init<const torch::Tensor &, const torch::Tensor &>(),
+              py::arg("in_aabb_mins"),
+              py::arg("in_aabb_maxs"),
+              "Construct MeshBVH from Triangle AABBs.")
+         .def("get_self_intersection", &MeshBVH::get_self_intersection,
+              py::arg("vertices"),
+              py::arg("triangles"),
+              "Find all self-intersecting triangle pairs")
+         .def("is_self_intersection", &MeshBVH::is_self_intersection,
+              py::arg("vertices"),
+              py::arg("triangles"),
+              "Check if there are any self-intersecting triangle pairs")
+         .def("get_ray_intersection", [](MeshBVH &self, const torch::Tensor &ray_origins, const torch::Tensor &ray_dirs, const torch::Tensor &vertices, const torch::Tensor &triangles, bool return_distance) -> py::object
+              {
+                 auto result = self.get_ray_intersection(ray_origins, ray_dirs, vertices, triangles, return_distance);
+                 if (return_distance) {
+                     return py::cast(result);
+                 } else {
+                     return py::cast(std::make_tuple(std::get<0>(result), std::get<1>(result), std::get<2>(result)));
+                 } }, py::arg("ray_origins"), py::arg("ray_dirs"), py::arg("vertices"), py::arg("triangles"), py::arg("return_distance") = false, "Find all ray-triangle intersections. Returns (ray_ids, triangle_ids, intersect_points, [distances])")
+         .def("query_point", &MeshBVH::query_point, py::arg("query_points"), py::arg("vertices"), py::arg("triangles"), py::arg("return_sdf") = false, py::arg("return_prj_pts") = true, py::arg("sign_mode") = 0, "Find the closest triangle to each query point. Returns (query_ids, triangle_ids, projected_points, distances)")
+         .def("build_winding_data", &MeshBVH::build_winding_data, py::arg("vertices"), py::arg("triangles"), "Computes the hierarchical area-weighted constants for Fast Winding Number queries.");
+}

{conquer3d-0.2.5 → conquer3d-0.2.7}/conquer3d/csrc/binds/data_structure/triangle_mesh.cpp

@@ -50,11 +50,33 @@ void TriangleMesh::compute_triangle_normals()
     this->triangle_normals = torch::empty({static_cast<int64_t>(this->num_triangles), 3}, torch::dtype(torch::kFloat32).device(this->vertices.device()));
     triangle_mesh::compute_triangle_normals(
         this->num_triangles,
-        reinterpret_cast<float3 *>(this->vertices.data_ptr<float>()),
-        reinterpret_cast<int3 *>(this->triangles.data_ptr<int>()),
+        reinterpret_cast<const float3 *>(this->vertices.data_ptr<float>()),
+        reinterpret_cast<const int3 *>(this->triangles.data_ptr<int>()),
         reinterpret_cast<float3 *>(this->triangle_normals.data_ptr<float>()));
 }
 
+void TriangleMesh::compute_vertex_normals()
+{
+    uint32_t num_vertices = this->vertices.size(0);
+    this->vertex_normals = torch::zeros({static_cast<int64_t>(num_vertices), 3}, torch::dtype(torch::kFloat32).device(this->vertices.device()));
+    
+    triangle_mesh::compute_vertex_normals(
+        num_vertices,
+        this->num_triangles,
+        reinterpret_cast<const int3 *>(this->triangles.data_ptr<int>()),
+        reinterpret_cast<const float3 *>(this->get_triangle_normals().data_ptr<float>()),
+        reinterpret_cast<float3 *>(this->vertex_normals.data_ptr<float>()));
+}
+
+torch::Tensor TriangleMesh::get_vertex_normals()
+{
+    if (!this->vertex_normals.defined())
+    {
+        this->compute_vertex_normals();
+    }
+    return this->vertex_normals;
+}
+
 void TriangleMesh::compute_triangle_areas()
 {
     this->triangle_areas = torch::empty({static_cast<int64_t>(this->num_triangles)}, torch::dtype(torch::kFloat32).device(this->vertices.device()));
@@ -97,6 +119,21 @@ bool TriangleMesh::is_self_intersection()
     return this->bvh.value().is_self_intersection(this->vertices, this->triangles);
 }
 
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> TriangleMesh::query_points(
+    const torch::Tensor &query_pts,
+    bool return_sdf,
+    bool return_prj_pts,
+    int sign_mode,
+    int distance_mode)
+{
+    if (distance_mode == 0) {
+        this->build_bvh();
+        return this->bvh.value().query_point(query_pts, this->vertices, this->triangles, return_sdf, return_prj_pts, sign_mode);
+    } else {
+        throw std::runtime_error("distance_mode != 0 not implemented yet");
+    }
+}
+
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> TriangleMesh::get_ray_intersection(
     const torch::Tensor &ray_origins,
     const torch::Tensor &ray_dirs,
@@ -106,6 +143,69 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> TriangleM
     return this->bvh.value().get_ray_intersection(ray_origins, ray_dirs, this->vertices, this->triangles, return_distance);
 }
 
+std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>, std::optional<torch::Tensor>> TriangleMesh::sample_points(
+    int num_points, bool uniform, bool return_normals, bool return_colors, bool use_triangle_normal)
+{
+    if (this->num_triangles == 0) {
+        throw std::runtime_error("Cannot sample points from an empty mesh.");
+    }
+
+    torch::Tensor tri_indices;
+    if (!uniform) {
+        tri_indices = torch::randint(0, this->num_triangles, {num_points}, torch::TensorOptions().device(this->vertices.device()).dtype(torch::kInt64));
+    } else {
+        torch::Tensor areas = this->get_triangle_areas();
+        tri_indices = torch::multinomial(areas, num_points, true);
+    }
+
+    torch::Tensor r1_r2 = torch::rand({num_points, 2}, torch::TensorOptions().device(this->vertices.device()).dtype(torch::kFloat32));
+    torch::Tensor out_points = torch::empty({num_points, 3}, torch::TensorOptions().device(this->vertices.device()).dtype(torch::kFloat32));
+
+    torch::Tensor out_normals, out_colors;
+    const float3* d_vertex_normals = nullptr;
+    const float3* d_triangle_normals = nullptr;
+    const float3* d_vertex_colors = nullptr;
+    float3* d_out_normals = nullptr;
+    float3* d_out_colors = nullptr;
+
+    if (return_normals) {
+        out_normals = torch::empty({num_points, 3}, torch::TensorOptions().device(this->vertices.device()).dtype(torch::kFloat32));
+        d_out_normals = reinterpret_cast<float3 *>(out_normals.data_ptr<float>());
+        if (use_triangle_normal) {
+            d_triangle_normals = reinterpret_cast<const float3 *>(this->get_triangle_normals().data_ptr<float>());
+        } else {
+            d_vertex_normals = reinterpret_cast<const float3 *>(this->get_vertex_normals().data_ptr<float>());
+        }
+    }
+
+    if (return_colors) {
+        if (!this->vertex_colors.defined()) {
+            throw std::runtime_error("Cannot sample colors because vertex_colors is not defined.");
+        }
+        out_colors = torch::empty({num_points, 3}, torch::TensorOptions().device(this->vertices.device()).dtype(torch::kFloat32));
+        d_out_colors = reinterpret_cast<float3 *>(out_colors.data_ptr<float>());
+        d_vertex_colors = reinterpret_cast<const float3 *>(this->vertex_colors.data_ptr<float>());
+    }
+
+    triangle_mesh::sample_points_triangle_mesh(
+        num_points,
+        reinterpret_cast<const float3 *>(this->vertices.data_ptr<float>()),
+        reinterpret_cast<const int3 *>(this->triangles.data_ptr<int>()),
+        reinterpret_cast<const int64_t *>(tri_indices.data_ptr<int64_t>()),
+        reinterpret_cast<const float2 *>(r1_r2.data_ptr<float>()),
+        d_vertex_normals,
+        d_triangle_normals,
+        d_vertex_colors,
+        reinterpret_cast<float3 *>(out_points.data_ptr<float>()),
+        d_out_normals,
+        d_out_colors);
+
+    std::optional<torch::Tensor> opt_normals = return_normals ? std::make_optional(out_normals) : std::nullopt;
+    std::optional<torch::Tensor> opt_colors = return_colors ? std::make_optional(out_colors) : std::nullopt;
+
+    return std::make_tuple(out_points, tri_indices, opt_normals, opt_colors);
+}
+
 torch::Tensor TriangleMesh::get_triangle_areas()
 {
     if (!this->triangle_areas.defined())
@@ -338,6 +438,13 @@ void bind_ds_triangle_mesh(py::module_ &m)
              py::arg("ray_dirs"),
              py::arg("return_distance") = false,
              "Find all ray-triangle intersections. Returns (ray_ids, triangle_ids, intersect_points, [distances])")
+        .def("query_points", &TriangleMesh::query_points,
+             py::arg("query_pts"),
+             py::arg("return_sdf") = false,
+             py::arg("return_prj_pts") = true,
+             py::arg("sign_mode") = 0,
+             py::arg("distance_mode") = 0,
+             "Query points against the mesh. Returns (query_ids, triangle_ids, projected_points, distances)")
         .def_property_readonly("edges", &TriangleMesh::get_edges)
         .def_property_readonly("edge_to_triangle_offsets", &TriangleMesh::get_edge_to_triangle_offsets)
         .def_property_readonly("edge_to_triangle_counts", &TriangleMesh::get_edge_to_triangle_counts)
@@ -356,6 +463,8 @@ void bind_ds_triangle_mesh(py::module_ &m)
         .def("is_manifold", &TriangleMesh::is_manifold, py::arg("allow_boundary_edge") = true)
         .def("get_non_manifold_vertices", &TriangleMesh::get_non_manifold_vertices)
         .def("remove_triangles_by_mask", &TriangleMesh::remove_triangles_by_mask, py::arg("keep_mask"))
+        .def("sample_points", &TriangleMesh::sample_points, py::arg("num_points"), py::arg("uniform") = false, py::arg("return_normals") = false, py::arg("return_colors") = false, py::arg("use_triangle_normal") = true, "Sample points on the mesh")
+        .def("compute_vertex_normals", &TriangleMesh::compute_vertex_normals)
         .def_property_readonly("euler_characteristic", &TriangleMesh::get_euler_characteristic)
         .def_property_readonly("genus", &TriangleMesh::get_genus);
 }