Migrate kd-tree operations to numpy arrays with RowMajor storage fix

pymomentum/array_utility/default_parameter_set.h

@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+namespace pymomentum {
+
+// Enum for controlling default behavior when converting empty arrays/tensors to ParameterSet.
+// If the user passes an empty tensor/array for a parameter set, what kind of value to return.
+// This is different for different cases: sometimes we should include all parameters,
+// sometimes none, and sometimes no reasonable default is possible.
+// clang-format off
+enum class DefaultParameterSet { AllOnes, AllZeros, NoDefault };
+// clang-format on
+
+} // namespace pymomentum

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_{};

pymomentum/backend/utils.py

@@ -140,7 +140,7 @@ def __init__(self, character: pym_geometry.Character, device: str = "cpu") -> No
 
         # Parameter transform properties
         self.param_transform: torch.Tensor = torch.tensor(
-            character.parameter_transform.transform.numpy(),
+            character.parameter_transform.transform,
             dtype=self.dtype,
             device=self._device,
         )
@@ -150,14 +150,9 @@ def __init__(self, character: pym_geometry.Character, device: str = "cpu") -> No
         )
 
         # Derived parameter counts from the boolean mask for backward compatibility
-        self.nr_position_params: int = int(
-            torch.logical_not(character.parameter_transform.scaling_parameters)
-            .sum()
-            .item()
-        )
-        self.nr_scaling_params: int = int(
-            character.parameter_transform.scaling_parameters.sum().item()
-        )
+        scaling_params = character.parameter_transform.scaling_parameters
+        self.nr_position_params: int = int((~scaling_params).sum())
+        self.nr_scaling_params: int = int(scaling_params.sum())
 
         # Mesh and skinning weights
         assert character.has_mesh

pymomentum/cmake/build_variables.bzl

@@ -31,6 +31,7 @@ tensor_utility_sources = [
 array_utility_public_headers = [
     "array_utility/array_utility.h",
     "array_utility/batch_accessor.h",
+    "array_utility/default_parameter_set.h",
     "array_utility/geometry_accessors.h",
 ]
 
@@ -118,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",
@@ -137,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",