Migrate MPPCA and finalize numpy migration with Python bindings (#967)

Summary:
Pull Request resolved: #967

Completed the migration from PyTorch tensors to numpy arrays for both MPPCA (Mixture of Probabilistic PCA) operations and kd-tree function bindings. This commit provides the Python bindings for both features.

Key changes:
- Implemented array_mppca for converting MPPCA models to numpy arrays
- Changed Mppca.to_tensors() to Mppca.to_arrays() returning numpy arrays instead of torch tensors
- Added Python bindings for find_closest_points, find_closest_points_with_normals, and find_closest_points_on_mesh using numpy arrays
- Removed dependency on ATen/PyTorch tensor headers from geometry module
- Updated test_posePrior to use numpy array indexing instead of tensor.select()

Note: This commit includes bindings for both the kd-tree functions (from the previous commit) and MPPCA, as they share common infrastructure changes (removing tensor includes and adding array includes).

Reviewed By: jeongseok-meta

Differential Revision: D89891113

Author: cdtwigg

pymomentum/cmake/build_variables.bzl

@@ -120,6 +120,7 @@ geometry_public_headers = [
     "geometry/array_blend_shape.h",
     "geometry/array_joint_parameters_to_positions.h",
     "geometry/array_kd_tree.h",
+    "geometry/array_mppca.h",
     "geometry/array_parameter_transform.h",
     "geometry/array_skeleton_state.h",
     "geometry/array_skinning.h",
@@ -140,6 +141,7 @@ geometry_sources = [
     "geometry/array_blend_shape.cpp",
     "geometry/array_joint_parameters_to_positions.cpp",
     "geometry/array_kd_tree.cpp",
+    "geometry/array_mppca.cpp",
     "geometry/array_parameter_transform.cpp",
     "geometry/array_skeleton_state.cpp",
     "geometry/array_skinning.cpp",

pymomentum/geometry/array_mppca.cpp

@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+#include "pymomentum/geometry/array_mppca.h"
+
+#include <momentum/common/exception.h>
+#include <momentum/math/constants.h>
+#include <momentum/math/mppca.h>
+
+#include <Eigen/Core>
+
+namespace py = pybind11;
+
+namespace pymomentum {
+
+std::tuple<
+    py::array_t<float>,
+    py::array_t<float>,
+    py::array_t<float>,
+    py::array_t<float>,
+    py::array_t<int>>
+mppcaToArrays(
+    const momentum::Mppca& mppca,
+    std::optional<const momentum::ParameterTransform*> paramTransform) {
+  const auto nMixtures = mppca.p;
+  const auto dimension = mppca.d;
+
+  // Create output arrays
+  py::array_t<float> pi_array(static_cast<py::ssize_t>(nMixtures));
+  py::array_t<float> mu_array(
+      {static_cast<py::ssize_t>(nMixtures), static_cast<py::ssize_t>(dimension)});
+
+  auto pi = pi_array.mutable_unchecked<1>();
+  auto mu = mu_array.mutable_unchecked<2>();
+
+  // Copy mu values
+  for (py::ssize_t iMix = 0; iMix < nMixtures; ++iMix) {
+    for (py::ssize_t d = 0; d < dimension; ++d) {
+      mu(iMix, d) = mppca.mu(iMix, d);
+    }
+  }
+
+  // Process each mixture component
+  MT_THROW_IF(mppca.Cinv.size() != nMixtures, "Invalid Mppca");
+
+  Eigen::VectorXf sigma_vec(nMixtures);
+  int W_rank = 0; // Will be determined from first mixture
+
+  // First pass: determine W rank from first mixture
+  {
+    Eigen::SelfAdjointEigenSolver<Eigen::MatrixXf> Cinv_eigs(mppca.Cinv[0]);
+    Eigen::VectorXf C_eigenvalues = Cinv_eigs.eigenvalues().cwiseInverse();
+    const float sigma2 = C_eigenvalues(C_eigenvalues.size() - 1);
+    C_eigenvalues.array() -= sigma2;
+
+    W_rank = C_eigenvalues.size();
+    for (Eigen::Index i = 0; i < C_eigenvalues.size(); ++i) {
+      if (C_eigenvalues(i) < 0.0001) {
+        W_rank = i;
+        break;
+      }
+    }
+  }
+
+  // Create W array with determined rank
+  py::array_t<float> W_array(
+      {static_cast<py::ssize_t>(nMixtures),
+       static_cast<py::ssize_t>(W_rank),
+       static_cast<py::ssize_t>(dimension)});
+  auto W = W_array.mutable_unchecked<3>();
+
+  // Second pass: fill in all values
+  for (Eigen::Index iMix = 0; iMix < nMixtures; ++iMix) {
+    Eigen::SelfAdjointEigenSolver<Eigen::MatrixXf> Cinv_eigs(mppca.Cinv[iMix]);
+
+    // Eigenvalues of the inverse are the inverse of the eigenvalues:
+    Eigen::VectorXf C_eigenvalues = Cinv_eigs.eigenvalues().cwiseInverse();
+
+    // Assume that it's not full rank and hence the last eigenvalue is sigma^2.
+    const float sigma2 = C_eigenvalues(C_eigenvalues.size() - 1);
+    assert(sigma2 >= 0);
+    sigma_vec[iMix] = std::sqrt(sigma2);
+
+    // (sigma^2*I + W^T*W) has eigenvalues (sigma^2 + lambda)
+    // where the lambda are the eigenvalues for W^T*W (which we want):
+    C_eigenvalues.array() -= sigma2;
+
+    // Fill W for this mixture
+    for (Eigen::Index jComponent = 0; jComponent < W_rank; ++jComponent) {
+      const float scale = std::sqrt(C_eigenvalues(jComponent));
+      for (Eigen::Index d = 0; d < dimension; ++d) {
+        W(iMix, jComponent, d) = scale * Cinv_eigs.eigenvectors()(d, jComponent);
+      }
+    }
+
+    const float C_logDeterminant = -Cinv_eigs.eigenvalues().array().log().sum();
+
+    // We have:
+    //   Rpre(c) = std::log(pi(c))
+    //       - 0.5 * C_logDeterminant
+    //       - 0.5 * static_cast<double>(d) * std::log(2.0 * PI));
+    // so std::log(pi(c)) = Rpre(c) + 0.5 * C_logDeterminant + 0.5 *
+    //      d * std::log(2.0 * PI));
+    const float log_pi = mppca.Rpre(iMix) + 0.5f * C_logDeterminant +
+        0.5f * static_cast<float>(mppca.d) * std::log(2.0 * momentum::pi<float>());
+    pi(iMix) = std::exp(log_pi);
+  }
+
+  // Create sigma array
+  py::array_t<float> sigma_array(static_cast<py::ssize_t>(nMixtures));
+  auto sigma = sigma_array.mutable_unchecked<1>();
+  for (py::ssize_t iMix = 0; iMix < nMixtures; ++iMix) {
+    sigma(iMix) = sigma_vec[iMix];
+  }
+
+  // Create parameter indices array
+  py::array_t<int> param_indices_array(static_cast<py::ssize_t>(dimension));
+  auto param_indices = param_indices_array.mutable_unchecked<1>();
+
+  for (py::ssize_t i = 0; i < dimension; ++i) {
+    param_indices(i) = -1; // Default to -1
+  }
+
+  if (paramTransform.has_value()) {
+    for (Eigen::Index i = 0; i < mppca.names.size() && i < dimension; ++i) {
+      auto paramIdx = (*paramTransform)->getParameterIdByName(mppca.names[i]);
+      if (paramIdx != momentum::kInvalidIndex) {
+        param_indices(i) = static_cast<int>(paramIdx);
+      }
+    }
+  }
+
+  return {pi_array, mu_array, W_array, sigma_array, param_indices_array};
+}
+
+} // namespace pymomentum

pymomentum/geometry/array_mppca.h

@@ -0,0 +1,37 @@
+/*
+ * 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
+
+#include <momentum/character/parameter_transform.h>
+#include <momentum/math/fwd.h>
+
+#include <pybind11/numpy.h>
+
+#include <optional>
+#include <tuple>
+
+namespace pymomentum {
+
+/// Convert an MPPCA model to numpy arrays.
+/// Returns (pi, mu, W, sigma, parameter_indices) where:
+/// - pi: [nMixtures] mixture weights
+/// - mu: [nMixtures, dimension] means
+/// - W: [nMixtures, rank, dimension] PCA basis vectors
+/// - sigma: [nMixtures] noise standard deviations
+/// - parameter_indices: [dimension] mapping to parameter transform indices (-1 if not found)
+std::tuple<
+    pybind11::array_t<float>,
+    pybind11::array_t<float>,
+    pybind11::array_t<float>,
+    pybind11::array_t<float>,
+    pybind11::array_t<int>>
+mppcaToArrays(
+    const momentum::Mppca& mppca,
+    std::optional<const momentum::ParameterTransform*> paramTransform);
+
+} // namespace pymomentum

pymomentum/geometry/geometry_pybind.cpp

@@ -8,6 +8,7 @@
 #include "pymomentum/geometry/array_blend_shape.h"
 #include "pymomentum/geometry/array_joint_parameters_to_positions.h"
 #include "pymomentum/geometry/array_kd_tree.h"
+#include "pymomentum/geometry/array_mppca.h"
 #include "pymomentum/geometry/array_parameter_transform.h"
 #include "pymomentum/geometry/array_skeleton_state.h"
 #include "pymomentum/geometry/array_vertex_normals.h"
@@ -22,11 +23,6 @@
 #include "pymomentum/geometry/skeleton_pybind.h"
 #include "pymomentum/geometry/skin_weights_pybind.h"
 
-// Keep tensor versions for functions without array equivalents yet
-#include "pymomentum/tensor_momentum/tensor_kd_tree.h"
-#include "pymomentum/tensor_momentum/tensor_mppca.h"
-#include "pymomentum/torch_bridge.h"
-
 #include <momentum/character/blend_shape.h>
 #include <momentum/character/character.h>
 #include <momentum/character/character_utility.h>
@@ -367,8 +363,8 @@ Each PPCA model is a Gaussian with mean mu and covariance (sigma^2*I + W*W^T).
       .def_readonly("n_dimension", &mm::Mppca::d, R"(The dimension of the parameter space.)")
       .def_readonly("names", &mm::Mppca::names, R"(The names of the parameters.)")
       .def(
-          "to_tensors",
-          &mppcaToTensors,
+          "to_arrays",
+          &mppcaToArrays,
           R"(Return the parameters defining the mixture of probabilistic PCA models.
 
 Each PPCA model a Gaussian N(mu, cov) where the covariance matrix is
@@ -377,16 +373,16 @@ Each PPCA model a Gaussian N(mu, cov) where the covariance matrix is
 Note that mu is a vector of length :meth:`dimension` and W is a matrix of dimension :meth:`dimension` x q
 where q is the dimensionality of the PCA subspace.
 
-The resulting tensors are as follows:
+The resulting arrays are as follows:
 
-* pi: a [n]-dimensional tensor containing the mixture weights.  It sums to 1.
-* mu: a [n x d]-dimensional tensor containing the mean pose for each mixture.
-* weights: a [n x d x q]-dimensional tensor containing the q vectors spanning the PCA space.
-* sigma: a [n]-dimensional tensor containing the uniform part of the covariance matrix.
-* param_idx: a [d]-dimensional tensor containing the indices of the parameters.
+* pi: a [n]-dimensional array containing the mixture weights.  It sums to 1.
+* mu: a [n x d]-dimensional array containing the mean pose for each mixture.
+* weights: a [n x q x d]-dimensional array containing the q vectors spanning the PCA space.
+* sigma: a [n]-dimensional array containing the uniform part of the covariance matrix.
+* param_idx: a [d]-dimensional array containing the indices of the parameters.
 
-:param parameter_transform: An optional parameter transform used to map the parameters; if not present, then the param_idx tensor will be empty.
-:return: an tuple (pi, mean, weights, sigma, param_idx) for the Probabilistic PCA model.)",
+:param parameter_transform: An optional parameter transform used to map the parameters; if not present, then the param_idx array will be empty.
+:return: a tuple (pi, mean, weights, sigma, param_idx) for the Probabilistic PCA model.)",
           py::arg("parameter_transform") = std::optional<const mm::ParameterTransform*>())
       .def("get_mixture", &getMppcaModel, py::arg("i_model"))
       .def_static(

pymomentum/geometry/momentum_geometry.h

@@ -10,7 +10,6 @@
 #include <momentum/character/character.h>
 #include <momentum/math/mppca.h>
 
-#include <ATen/ATen.h>
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>