Migrate kd-tree operations to numpy arrays with RowMajor storage fix (#970)

Summary:
Pull Request resolved: #970

Migrated find_closest_points, find_closest_points_with_normals, and find_closest_points_on_mesh from torch tensors to numpy arrays. This continues the broader migration of pymomentum geometry operations from PyTorch to numpy.

The key technical challenge was ensuring correct memory layout for the SimdKdTree library. Eigen matrices use column-major storage by default, but SimdKdTree expects row-major data where each point's coordinates are contiguous. Added a conditional toMatrix() method to VectorArrayAccessor that uses RowMajor storage for multi-dimensional points (Dim > 1) while preserving ColMajor for column vectors (Dim == 1) as required by Eigen.

Reviewed By: jeongseok-meta

Differential Revision: D89891110

Author: cdtwigg

pymomentum/array_utility/geometry_accessors.cpp

@@ -586,7 +586,32 @@ void VectorArrayAccessor<T, Dim>::set(
   }
 }
 
+template <typename T, int Dim>
+Eigen::Matrix<T, Eigen::Dynamic, Dim, (Dim > 1) ? Eigen::RowMajor : Eigen::ColMajor>
+VectorArrayAccessor<T, Dim>::toMatrix(const std::vector<py::ssize_t>& batchIndices) const {
+  const auto offset = computeOffset(batchIndices);
+
+  // Create matrix with nElements rows and Dim columns
+  // Use RowMajor storage for Dim > 1 to ensure each point's coordinates are contiguous in memory
+  // Use ColMajor (default) for Dim == 1 (column vectors) as required by Eigen
+  Eigen::Matrix<T, Eigen::Dynamic, Dim, (Dim > 1) ? Eigen::RowMajor : Eigen::ColMajor> result(
+      nElements_, Dim);
+
+  // Copy data row by row, handling arbitrary strides
+  for (py::ssize_t i = 0; i < nElements_; ++i) {
+    const auto elemOffset = offset + i * rowStride_;
+    for (int d = 0; d < Dim; ++d) {
+      result(i, d) = data_[elemOffset + d * colStride_];
+    }
+  }
+
+  return result;
+}
+
 // Explicit template instantiations for VectorArrayAccessor
+template class VectorArrayAccessor<int32_t, 1>;
+template class VectorArrayAccessor<float, 2>;
+template class VectorArrayAccessor<double, 2>;
 template class VectorArrayAccessor<float, 3>;
 template class VectorArrayAccessor<double, 3>;
 template class VectorArrayAccessor<int, 3>;

pymomentum/array_utility/geometry_accessors.h

@@ -364,6 +364,14 @@ class VectorArrayAccessor {
   // Set all vectors for the given batch indices (bulk operation).
   void set(const std::vector<py::ssize_t>& batchIndices, const std::vector<VectorType>& values);
 
+  // Convert to Eigen matrix format (nElements x Dim) with appropriate storage.
+  // This is useful for algorithms that expect Eigen::MatrixX3 or similar formats.
+  // Creates a new matrix with copied data, handling arbitrary strides.
+  // Uses RowMajor storage for Dim > 1 to ensure each point's coordinates are contiguous in memory.
+  // Uses ColMajor (default) for Dim == 1 (column vectors) as required by Eigen.
+  Eigen::Matrix<T, Eigen::Dynamic, Dim, (Dim > 1) ? Eigen::RowMajor : Eigen::ColMajor> toMatrix(
+      const std::vector<py::ssize_t>& batchIndices) const;
+
  private:
   T* data_;
   py::ssize_t nElements_{};
@@ -415,6 +423,50 @@ class IntVectorArrayAccessor {
     // Returns {min, max}. For empty views, returns {INT_MAX, INT_MIN}.
     [[nodiscard]] std::pair<int, int> minmax() const;
 
+    // Convert to Eigen matrix format with int values.
+    // This performs a single switch on dtype for the entire conversion,
+    // avoiding per-element overhead. Template parameter allows caller to
+    // specify the desired output matrix type (e.g., Eigen::MatrixX3i for
+    // column-major, or Eigen::Matrix<int, Dynamic, 3, RowMajor> for row-major).
+    template <typename MatrixType = Eigen::Matrix<int, Eigen::Dynamic, Dim>>
+    [[nodiscard]] MatrixType toMatrix() const {
+      MatrixType result(nElements_, Dim);
+
+      if (nElements_ == 0) {
+        return result;
+      }
+
+      // Helper lambda to copy data for a given source type
+      auto copyData = [&]<typename SourceT>() {
+        const auto elemStride = rowStride_ / static_cast<py::ssize_t>(sizeof(SourceT));
+        const auto compStride = colStride_ / static_cast<py::ssize_t>(sizeof(SourceT));
+        const auto* ptr = static_cast<const SourceT*>(data_);
+
+        for (py::ssize_t i = 0; i < nElements_; ++i) {
+          for (int d = 0; d < Dim; ++d) {
+            result(i, d) = static_cast<int>(ptr[i * elemStride + d * compStride]);
+          }
+        }
+      };
+
+      switch (dtype_) {
+        case SourceDtype::Int32:
+          copyData.template operator()<int32_t>();
+          break;
+        case SourceDtype::Int64:
+          copyData.template operator()<int64_t>();
+          break;
+        case SourceDtype::UInt32:
+          copyData.template operator()<uint32_t>();
+          break;
+        case SourceDtype::UInt64:
+          copyData.template operator()<uint64_t>();
+          break;
+      }
+
+      return result;
+    }
+
    private:
     friend class IntVectorArrayAccessor;
 

pymomentum/cmake/build_variables.bzl

@@ -119,6 +119,7 @@ tensor_ik_test_sources = [
 geometry_public_headers = [
     "geometry/array_blend_shape.h",
     "geometry/array_joint_parameters_to_positions.h",
+    "geometry/array_kd_tree.h",
     "geometry/array_parameter_transform.h",
     "geometry/array_skeleton_state.h",
     "geometry/array_skinning.h",
@@ -138,6 +139,7 @@ geometry_public_headers = [
 geometry_sources = [
     "geometry/array_blend_shape.cpp",
     "geometry/array_joint_parameters_to_positions.cpp",
+    "geometry/array_kd_tree.cpp",
     "geometry/array_parameter_transform.cpp",
     "geometry/array_skeleton_state.cpp",
     "geometry/array_skinning.cpp",