pymomentum-cpu 0.1.93.post0__cp312-cp312-macosx_14_0_arm64.whl

include/momentum/character_solver/trust_region_qr.h

@@ -0,0 +1,80 @@
+/*
+ * 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/mesh_state.h>
+#include <momentum/character/skeleton_state.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_solver_function.h>
+#include <momentum/math/online_householder_qr.h>
+#include <momentum/solver/solver.h>
+
+#include <Eigen/Core>
+
+namespace momentum {
+
+/// Trust Region with QR decomposition specific options
+struct TrustRegionQROptions : SolverOptions {
+  /// Trust region radius parameter for the Trust Region solver with QR decomposition.
+  float trustRegionRadius_ = 1.0f;
+
+  TrustRegionQROptions() = default;
+
+  /* implicit */ TrustRegionQROptions(const SolverOptions& baseOptions)
+      : SolverOptions(baseOptions) {
+    // Empty
+  }
+};
+
+template <typename T>
+class TrustRegionQRT : public SolverT<T> {
+ public:
+  TrustRegionQRT(const SolverOptions& options, SkeletonSolverFunctionT<T>* solverFun);
+
+  [[nodiscard]] std::string_view getName() const override;
+
+  void setOptions(const SolverOptions& options) final;
+
+  void iter() {
+    doIteration();
+    this->iteration_++;
+  }
+
+  void init() {
+    this->iteration_ = 0;
+    initializeSolver();
+  }
+
+  Eigen::VectorX<T> getP() {
+    return this->parameters_;
+  }
+
+ protected:
+  void doIteration() final;
+  void initializeSolver() final;
+
+ private:
+  ResizeableMatrix<T> jacobian_;
+  ResizeableMatrix<T> residual_;
+
+  Eigen::VectorX<T> gradientSub_;
+  Eigen::MatrixX<T> lambdaDiag_;
+  Eigen::VectorX<T> lambdaZero_;
+
+  OnlineHouseholderQR<T> qrSolver_;
+
+  SkeletonStateT<T> skeletonState_;
+  MeshStateT<T> meshState_;
+
+  T trustRegionRadius_;
+  T maxTrustRegionRadius_ = 10.0;
+  T curTrustRegionRadius_ = 1.0;
+  Eigen::MatrixX<T> Rmatrix_;
+};
+
+} // namespace momentum

include/momentum/character_solver/vertex_error_function.h

@@ -0,0 +1,155 @@
+/*
+ * 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_solver/fwd.h>
+#include <momentum/character_solver/plane_error_function.h>
+#include <momentum/character_solver/position_error_function.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/math/fwd.h>
+
+namespace momentum {
+
+template <typename T>
+struct VertexConstraintT {
+  int vertexIndex = -1;
+  T weight = 1;
+  Eigen::Vector3<T> targetPosition;
+  Eigen::Vector3<T> targetNormal;
+
+  template <typename T2>
+  VertexConstraintT<T2> cast() const {
+    return {
+        this->vertexIndex,
+        static_cast<T2>(this->weight),
+        this->targetPosition.template cast<T2>(),
+        this->targetNormal.template cast<T2>()};
+  }
+};
+
+enum class VertexConstraintType {
+  Position, // Target the vertex position
+  Plane, // point-to-plane distance using the target normal
+  Normal, // point-to-plane distance using the source (body) normal
+  SymmetricNormal, // Point-to-plane using a 50/50 mix of source and target normal
+};
+
+[[nodiscard]] std::string_view toString(VertexConstraintType type);
+
+template <typename T>
+class VertexErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  explicit VertexErrorFunctionT(
+      const Character& character,
+      VertexConstraintType type = VertexConstraintType::Position,
+      uint32_t maxThreads = 0);
+  virtual ~VertexErrorFunctionT() override;
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState) final;
+
+  double getGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  void addConstraint(
+      int vertexIndex,
+      T weight,
+      const Eigen::Vector3<T>& targetPosition,
+      const Eigen::Vector3<T>& targetNormal);
+  void clearConstraints();
+
+  [[nodiscard]] const std::vector<VertexConstraintT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+
+  static constexpr T kPositionWeight = PositionErrorFunctionT<T>::kLegacyWeight;
+  static constexpr T kPlaneWeight = PlaneErrorFunctionT<T>::kLegacyWeight;
+
+  [[nodiscard]] const Character* getCharacter() const override {
+    return &character_;
+  }
+
+  /// Override to indicate this function requires mesh state
+  [[nodiscard]] bool needsMesh() const override {
+    return true;
+  }
+
+ private:
+  double calculatePositionJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexConstraintT<T>& constr,
+      Ref<Eigen::MatrixX<T>> jac,
+      Ref<Eigen::VectorX<T>> res) const;
+
+  double calculateNormalJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexConstraintT<T>& constr,
+      T sourceNormalWeight,
+      T targetNormalWeight,
+      Ref<Eigen::MatrixX<T>> jac,
+      T& res) const;
+
+  double calculatePositionGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexConstraintT<T>& constr,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  double calculateNormalGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexConstraintT<T>& constr,
+      T sourceNormalWeight,
+      T targetNormalWeight,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  // Utility function used now in calculateNormalJacobian and calculatePositionGradient
+  // to calculate derivatives with respect to position in world space (considering skinning)
+  void calculateDWorldPos(
+      const SkeletonStateT<T>& state,
+      const VertexConstraintT<T>& constr,
+      const Eigen::Vector3<T>& d_restPos,
+      Eigen::Vector3<T>& d_worldPos) const;
+
+  std::pair<T, T> computeNormalWeights() const;
+
+  const Character& character_;
+
+  std::vector<VertexConstraintT<T>> constraints_;
+
+  const VertexConstraintType constraintType_;
+
+  uint32_t maxThreads_;
+};
+
+} // namespace momentum

include/momentum/character_solver/vertex_projection_error_function.h

@@ -0,0 +1,117 @@
+/*
+ * 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_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/math/fwd.h>
+
+namespace momentum {
+
+template <typename T>
+struct VertexProjectionConstraintT {
+  int vertexIndex = -1;
+  T weight = 1;
+  Eigen::Vector2<T> targetPosition;
+  Eigen::Matrix<T, 3, 4> projection; // Projection matrix
+
+  template <typename T2>
+  VertexProjectionConstraintT<T2> cast() const {
+    return {
+        this->vertexIndex,
+        (T)this->weight,
+        this->targetPosition.template cast<T2>(),
+    };
+  }
+};
+
+template <typename T>
+class VertexProjectionErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  explicit VertexProjectionErrorFunctionT(const Character& character, uint32_t maxThreads = 0);
+  ~VertexProjectionErrorFunctionT() override;
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState) final;
+
+  double getGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  void addConstraint(
+      int vertexIndex,
+      T weight,
+      const Eigen::Vector2<T>& targetPosition,
+      const Eigen::Matrix<T, 3, 4>& projection);
+  void clearConstraints();
+
+  [[nodiscard]] const std::vector<VertexProjectionConstraintT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+  [[nodiscard]] const Character* getCharacter() const override {
+    return &character_;
+  }
+
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+
+  /// Override to indicate this function requires mesh state
+  [[nodiscard]] bool needsMesh() const override {
+    return true;
+  }
+
+ private:
+  double calculateJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexProjectionConstraintT<T>& constr,
+      Ref<Eigen::MatrixX<T>> jac,
+      Ref<Eigen::VectorX<T>> res) const;
+
+  double calculateGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexProjectionConstraintT<T>& constr,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  // Utility function used now in calculateJacobian and calculateGradient
+  // to calculate derivatives with respect to position in world space (considering skinning)
+  void calculateDWorldPos(
+      const SkeletonStateT<T>& state,
+      const VertexProjectionConstraintT<T>& constr,
+      const Eigen::Vector3<T>& d_restPos,
+      Eigen::Vector3<T>& d_worldPos) const;
+
+  const Character& character_;
+
+  std::vector<VertexProjectionConstraintT<T>> constraints_;
+
+  uint32_t maxThreads_;
+
+  T _nearClip = 1.0f;
+};
+
+} // namespace momentum

include/momentum/character_solver/vertex_vertex_distance_error_function.h

@@ -0,0 +1,147 @@
+/*
+ * 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_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/math/fwd.h>
+
+#include <Eigen/Dense>
+#include <memory>
+#include <vector>
+
+namespace momentum {
+
+/// Constraint for vertex-to-vertex distance errors
+template <typename T>
+struct VertexVertexDistanceConstraintT {
+  int vertexIndex1 = -1; ///< First vertex index
+  int vertexIndex2 = -1; ///< Second vertex index
+  T weight = 1; ///< Constraint weight
+  T targetDistance = 0; ///< Desired distance between the two vertices
+
+  template <typename T2>
+  VertexVertexDistanceConstraintT<T2> cast() const {
+    return {
+        this->vertexIndex1,
+        this->vertexIndex2,
+        static_cast<T2>(this->weight),
+        static_cast<T2>(this->targetDistance)};
+  }
+};
+
+/// Error function for vertex-to-vertex distance constraints
+template <typename T>
+class VertexVertexDistanceErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  explicit VertexVertexDistanceErrorFunctionT(const Character& character);
+  ~VertexVertexDistanceErrorFunctionT() override;
+
+  VertexVertexDistanceErrorFunctionT(const VertexVertexDistanceErrorFunctionT& other) = delete;
+  VertexVertexDistanceErrorFunctionT(VertexVertexDistanceErrorFunctionT&& other) noexcept = delete;
+  VertexVertexDistanceErrorFunctionT& operator=(const VertexVertexDistanceErrorFunctionT& other) =
+      delete;
+  VertexVertexDistanceErrorFunctionT& operator=(VertexVertexDistanceErrorFunctionT&& other) =
+      delete;
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState) final;
+
+  double getGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  /// Add a vertex-to-vertex distance constraint
+  void addConstraint(int vertexIndex1, int vertexIndex2, T weight, T targetDistance);
+
+  /// Clear all constraints
+  void clearConstraints();
+
+  /// Get all constraints
+  [[nodiscard]] const std::vector<VertexVertexDistanceConstraintT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+  /// Get the number of constraints
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+
+  /// Get the character
+  [[nodiscard]] const Character* getCharacter() const override {
+    return &character_;
+  }
+
+  /// Override to indicate this function requires mesh state
+  [[nodiscard]] bool needsMesh() const override {
+    return true;
+  }
+
+ private:
+  /// Calculate jacobian for a distance constraint
+  double calculateJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexVertexDistanceConstraintT<T>& constraint,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      T& residual) const;
+
+  /// Calculate gradient for a distance constraint
+  double calculateGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const VertexVertexDistanceConstraintT<T>& constraint,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  /// Calculate gradient contribution from a single vertex
+  void calculateVertexGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      int vertexIndex,
+      const Eigen::Vector3<T>& gradientDirection,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  /// Calculate jacobian contribution from a single vertex
+  void calculateVertexJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      int vertexIndex,
+      const Eigen::Vector3<T>& jacobianDirection,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian) const;
+
+  /// Calculate world space position derivative for blend shape parameters
+  void calculateDWorldPos(
+      const SkeletonStateT<T>& state,
+      int vertexIndex,
+      const Eigen::Vector3<T>& d_restPos,
+      Eigen::Vector3<T>& d_worldPos) const;
+
+  const Character& character_;
+
+  std::vector<VertexVertexDistanceConstraintT<T>> constraints_;
+};
+
+} // namespace momentum

include/momentum/common/aligned.h

@@ -0,0 +1,155 @@
+/*
+ * 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/common/exception.h>
+
+#include <limits>
+#include <memory>
+
+namespace momentum {
+
+#if defined(_WIN32)
+
+[[nodiscard]] inline void* aligned_malloc(size_t size, size_t align) {
+  return _aligned_malloc(size, align);
+}
+
+inline void aligned_free(void* ptr) {
+  return _aligned_free(ptr);
+}
+
+#elif (defined(__IPHONE_OS_VERSION_MIN_REQUIRED) && __IPHONE_OS_VERSION_MIN_REQUIRED < 130000) || \
+    (defined(__ANDROID_API__) && __ANDROID_API__ < 28)
+
+[[nodiscard]] inline void* aligned_malloc(size_t size, size_t align) {
+  void* result = nullptr;
+  posix_memalign(&result, align, size);
+  return result;
+}
+
+inline void aligned_free(void* ptr) {
+  return std::free(ptr);
+}
+
+#else
+
+[[nodiscard]] inline void* aligned_malloc(size_t size, size_t align) {
+  return std::aligned_alloc(align, size);
+}
+
+inline void aligned_free(void* ptr) {
+  return std::free(ptr);
+}
+
+#endif
+
+inline constexpr std::size_t roundUpToAlignment(std::size_t value, std::size_t alignment) {
+  MT_THROW_IF_T(alignment == 0, std::invalid_argument, "Alignment must be non-zero");
+  return ((value + alignment - 1) / alignment) * alignment;
+}
+
+/// Allocates a block of memory that can hold `n` elements of type `T` with the specified alignment.
+///
+/// This function is intended to be used in the `AlignedAllocator::allocate()` method and throws
+/// exceptions as `std::allocator<T>::allocate` does.
+template <typename T, std::size_t Alignment = alignof(T)>
+[[nodiscard]] T* alignedAlloc(std::size_t n) {
+  MT_THROW_IF_T(std::numeric_limits<std::size_t>::max() / sizeof(T) < n, std::bad_array_new_length);
+
+  const std::size_t size = roundUpToAlignment(n * sizeof(T), Alignment);
+  void* ptr = aligned_malloc(size, Alignment);
+
+  MT_THROW_IF_T(ptr == nullptr, std::bad_alloc);
+
+  return static_cast<T*>(ptr);
+}
+
+/// Custom deleter for aligned memory
+struct AlignedDeleter {
+  void operator()(void* ptr) const {
+    aligned_free(ptr);
+  }
+};
+
+/// Creates a std::unique_ptr for aligned memory.
+template <typename T, std::size_t Alignment = alignof(T), class Deleter = AlignedDeleter>
+[[nodiscard]] std::unique_ptr<T, Deleter> makeAlignedUnique(
+    std::size_t n,
+    Deleter deleter = Deleter()) {
+  return std::unique_ptr<T, Deleter>(alignedAlloc<T, Alignment>(n), std::move(deleter));
+}
+
+/// Creates a std::shared_ptr for aligned memory.
+template <typename T, std::size_t Alignment = alignof(T), class Deleter = AlignedDeleter>
+[[nodiscard]] std::shared_ptr<T> makeAlignedShared(std::size_t n, Deleter deleter = Deleter()) {
+  return std::shared_ptr<T>(alignedAlloc<T, Alignment>(n), std::move(deleter));
+}
+
+/// An allocator that aligns memory blocks according to a specified alignment.
+/// The allocator is compatible with `std::allocator` and can be used in
+/// place of `std::allocator` in STL containers.
+///
+/// @tparam T The type of elements that the allocator will allocate.
+/// @tparam Alignment The alignment for the allocated memory blocks.
+template <class T, std::size_t Alignment>
+class AlignedAllocator {
+ public:
+  using value_type = T;
+
+  AlignedAllocator() noexcept = default;
+
+  template <class U>
+  explicit AlignedAllocator(const AlignedAllocator<U, Alignment>& /*other*/) noexcept {
+    // Empty
+  }
+
+  /// Allocates a block of memory that can hold `n` elements of type `T`.
+  ///
+  /// @param[in] n The number of elements to allocate space for.
+  /// @return A pointer to the first byte of the allocated memory block.
+  [[nodiscard]] T* allocate(std::size_t n) {
+    return alignedAlloc<T, Alignment>(n);
+  }
+
+  /// Deallocates a block of memory that was previously allocated by `allocate`.
+  ///
+  /// @param[in] ptr A pointer to the first byte of the memory block to deallocate.
+  /// @param size The size of the memory block to deallocate. This parameter is not used, but it is
+  /// included to maintain compatibility with `std::allocator`.
+  void deallocate(T* ptr, std::size_t /*n*/) noexcept {
+    aligned_free(ptr);
+  }
+
+  template <class U>
+  struct rebind {
+    using other = AlignedAllocator<U, Alignment>;
+  };
+};
+
+/// Checks if storage allocated from `lhs` can be deallocated from `rhs`, and vice versa.
+///
+/// This always returns true for stateless allocators.
+template <class T, class U, std::size_t Alignment>
+[[nodiscard]] bool operator==(
+    const AlignedAllocator<T, Alignment>& /*lhs*/,
+    const AlignedAllocator<U, Alignment>& /*rhs*/) {
+  return true;
+}
+
+/// Checks if storage allocated from `lhs` cannot be deallocated from `rhs`, and vice versa.
+///
+/// This always returns false for stateless allocators.
+template <class T, class U, std::size_t Alignment>
+[[nodiscard]] bool operator!=(
+    const AlignedAllocator<T, Alignment>& /*lhs*/,
+    const AlignedAllocator<U, Alignment>& /*rhs*/) {
+  return false;
+}
+
+} // namespace momentum

include/momentum/common/checks.h

@@ -0,0 +1,27 @@
+/*
+ * 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
+
+#if defined(MOMENTUM_WITH_XR_LOGGER)
+
+#include <logging/Checks.h>
+
+#define MT_CHECK(condition, ...) XR_CHECK(condition, __VA_ARGS__)
+#define MT_CHECK_LT(val1, val2, ...) XR_CHECK_DETAIL_OP1(val1, val2, <, ##__VA_ARGS__)
+#define MT_CHECK_NOTNULL(ptr, ...) XR_CHECK(ptr != nullptr, ##__VA_ARGS__)
+
+#else
+
+#include <cassert>
+
+// TODO: Support asserts with messages as XR_CHECK does
+#define MT_CHECK(condition, ...) assert(condition)
+#define MT_CHECK_LT(val1, val2, ...) assert(val1 < val2)
+#define MT_CHECK_NOTNULL(ptr, ...) assert(ptr != nullptr)
+
+#endif