Skip to content

Commit e43e716

Browse files
cdtwiggmeta-codesync[bot]
authored andcommitted
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
1 parent fa30743 commit e43e716

File tree

8 files changed

+1050
-71
lines changed

8 files changed

+1050
-71
lines changed

pymomentum/array_utility/geometry_accessors.cpp

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -586,7 +586,32 @@ void VectorArrayAccessor<T, Dim>::set(
586586
}
587587
}
588588

589+
template <typename T, int Dim>
590+
Eigen::Matrix<T, Eigen::Dynamic, Dim, (Dim > 1) ? Eigen::RowMajor : Eigen::ColMajor>
591+
VectorArrayAccessor<T, Dim>::toMatrix(const std::vector<py::ssize_t>& batchIndices) const {
592+
const auto offset = computeOffset(batchIndices);
593+
594+
// Create matrix with nElements rows and Dim columns
595+
// Use RowMajor storage for Dim > 1 to ensure each point's coordinates are contiguous in memory
596+
// Use ColMajor (default) for Dim == 1 (column vectors) as required by Eigen
597+
Eigen::Matrix<T, Eigen::Dynamic, Dim, (Dim > 1) ? Eigen::RowMajor : Eigen::ColMajor> result(
598+
nElements_, Dim);
599+
600+
// Copy data row by row, handling arbitrary strides
601+
for (py::ssize_t i = 0; i < nElements_; ++i) {
602+
const auto elemOffset = offset + i * rowStride_;
603+
for (int d = 0; d < Dim; ++d) {
604+
result(i, d) = data_[elemOffset + d * colStride_];
605+
}
606+
}
607+
608+
return result;
609+
}
610+
589611
// Explicit template instantiations for VectorArrayAccessor
612+
template class VectorArrayAccessor<int32_t, 1>;
613+
template class VectorArrayAccessor<float, 2>;
614+
template class VectorArrayAccessor<double, 2>;
590615
template class VectorArrayAccessor<float, 3>;
591616
template class VectorArrayAccessor<double, 3>;
592617
template class VectorArrayAccessor<int32_t, 3>;

pymomentum/array_utility/geometry_accessors.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,14 @@ class VectorArrayAccessor {
362362
// Set all vectors for the given batch indices (bulk operation).
363363
void set(const std::vector<py::ssize_t>& batchIndices, const std::vector<VectorType>& values);
364364

365+
// Convert to Eigen matrix format (nElements x Dim) with appropriate storage.
366+
// This is useful for algorithms that expect Eigen::MatrixX3 or similar formats.
367+
// Creates a new matrix with copied data, handling arbitrary strides.
368+
// Uses RowMajor storage for Dim > 1 to ensure each point's coordinates are contiguous in memory.
369+
// Uses ColMajor (default) for Dim == 1 (column vectors) as required by Eigen.
370+
Eigen::Matrix<T, Eigen::Dynamic, Dim, (Dim > 1) ? Eigen::RowMajor : Eigen::ColMajor> toMatrix(
371+
const std::vector<py::ssize_t>& batchIndices) const;
372+
365373
private:
366374
T* data_;
367375
py::ssize_t nElements_{};

pymomentum/cmake/build_variables.bzl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ tensor_ik_test_sources = [
119119
geometry_public_headers = [
120120
"geometry/array_blend_shape.h",
121121
"geometry/array_joint_parameters_to_positions.h",
122+
"geometry/array_kd_tree.h",
122123
"geometry/array_parameter_transform.h",
123124
"geometry/array_skeleton_state.h",
124125
"geometry/array_skinning.h",
@@ -138,6 +139,7 @@ geometry_public_headers = [
138139
geometry_sources = [
139140
"geometry/array_blend_shape.cpp",
140141
"geometry/array_joint_parameters_to_positions.cpp",
142+
"geometry/array_kd_tree.cpp",
141143
"geometry/array_parameter_transform.cpp",
142144
"geometry/array_skeleton_state.cpp",
143145
"geometry/array_skinning.cpp",

0 commit comments

Comments
 (0)