pymomentum-cpu 0.0.0__cp312-cp312-manylinux_2_39_x86_64.whl

include/momentum/character_solver/skinning_weight_iterator.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/fwd.h>
+#include <momentum/character/skin_weights.h>
+#include <momentum/math/fwd.h>
+
+namespace momentum {
+
+template <typename T>
+struct BoneWeightT {
+  size_t parentBone{};
+  float weight{};
+  Eigen::Vector3<T> weightedWorldSpacePoint;
+
+  bool operator<(const BoneWeightT<T>& rhs) const {
+    return parentBone < rhs.parentBone;
+  }
+
+  bool operator>(const BoneWeightT<T>& rhs) const {
+    return parentBone > rhs.parentBone;
+  }
+
+  BoneWeightT<T>& operator+=(const BoneWeightT<T>& rhs) {
+    this->weight += rhs.weight;
+    this->weightedWorldSpacePoint += rhs.weightedWorldSpacePoint;
+    return *this;
+  }
+};
+
+// When we compute the derivatives of skinned points, we usually
+// end up with a loop like this:
+//   for each skinned point i
+//     for each skinning bone and weight B_ij, w_ij
+//       for each ancestor k
+//          compute the weighted change in the skinned point wrt the ancestor
+// The redundancy here is that most points are skinned to a set of bones in the
+// same hierarchy (for example, two parts of the arm) and so it's redundant to
+// walk up the whole hierarchy from scratch for each skinned bone.
+//
+// This class instead does a single pass up the tree from leaf-most to root-most
+// skinned bones, accumulating skinning weights along the way.  By packaging it
+// in this class we hide the complexity and simplify the jacobian/gradient
+// calculation.
+template <typename T>
+class SkinningWeightIteratorT {
+ public:
+  SkinningWeightIteratorT(
+      const Character& character,
+      const MeshT<T>& restMesh,
+      const SkeletonStateT<T>& skelState,
+      int vertexIndex);
+
+  SkinningWeightIteratorT(
+      const Character& character,
+      const SkinnedLocator& locator,
+      const Eigen::Vector3<T>& locatorPosition,
+      const SkeletonStateT<T>& skelState);
+
+  [[nodiscard]] bool finished() const;
+
+  // Returns the tuple <parent bone index, bone weight, vertex position in world space wrt the
+  // current bone>
+  std::tuple<size_t, T, Eigen::Vector3<T>> next();
+
+ private:
+  void checkInvariants();
+
+  std::array<BoneWeightT<T>, kMaxSkinJoints> boneWeights;
+  int nBoneWeights;
+  const Character& character;
+};
+
+} // namespace momentum

include/momentum/character_solver/state_error_function.h

@@ -0,0 +1,90 @@
+/*
+ * 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/skeleton_state.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+
+namespace momentum {
+
+template <typename T>
+class StateErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  StateErrorFunctionT(const Skeleton& skel, const ParameterTransform& pt);
+  explicit StateErrorFunctionT(const Character& character);
+
+  [[nodiscard]] double getError(const ModelParametersT<T>& params, const SkeletonStateT<T>& state)
+      final;
+
+  double getGradient(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  void reset();
+  void setTargetState(const SkeletonStateT<T>* target);
+  void setTargetState(const SkeletonStateT<T>& target);
+  void setTargetState(TransformListT<T> target);
+  void setTargetWeight(const Eigen::VectorX<T>& weights);
+  void setTargetWeights(const Eigen::VectorX<T>& posWeight, const Eigen::VectorX<T>& rotWeight);
+
+  /// Set the target position weights for each joint.
+  /// @param posWeight Per-joint position weights.
+  void setPositionTargetWeights(const Eigen::VectorX<T>& posWeight);
+
+  /// Set the target rotation weights for each joint.
+  /// @param rotWeight Per-joint rotation weights.
+  void setRotationTargetWeights(const Eigen::VectorX<T>& rotWeight);
+
+  void setWeights(const float posWeight, const float rotationWeight) {
+    posWgt_ = posWeight;
+    rotWgt_ = rotationWeight;
+  }
+
+  [[nodiscard]] const TransformListT<T>& getTargetState() const {
+    return this->targetState_;
+  }
+
+  [[nodiscard]] const Eigen::VectorX<T>& getPositionWeights() const {
+    return targetPositionWeights_;
+  }
+  [[nodiscard]] const Eigen::VectorX<T>& getRotationWeights() const {
+    return targetRotationWeights_;
+  }
+  [[nodiscard]] const T& getPositionWeight() const {
+    return posWgt_;
+  }
+  [[nodiscard]] const T& getRotationWeight() const {
+    return rotWgt_;
+  }
+
+ private:
+  TransformListT<T> targetState_;
+  Eigen::VectorX<T> targetPositionWeights_;
+  Eigen::VectorX<T> targetRotationWeights_;
+
+  T posWgt_;
+  T rotWgt_;
+
+ public:
+  // weights for the error functions
+  static constexpr T kPositionWeight = 1e-3f;
+  static constexpr T kOrientationWeight = 1e+0f;
+};
+
+} // namespace momentum

include/momentum/character_solver/transform_pose.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/character.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/math/random.h>
+#include <momentum/math/transform.h>
+#include <momentum/math/types.h>
+#include <momentum/solver/fwd.h>
+
+namespace momentum {
+
+/// Function which computes a new set of model parameters such that the character pose is a rigidly
+/// transformed version of the original pose.  While there is technically a closed form solution
+/// for any given skeleton, this is complicated in momentum because different
+/// characters attach the rigid parameters to different joints, so a fully general solution uses IK.
+/// However, getting it right requires dealing with local minima issues in Euler angles and other
+/// challenges.  So centralizing this functionality in a single place is useful.
+///
+/// @param[in] ensureContinuousOutput If true, the output model parameters will be solved
+/// sequentially so that the resulting output is continuous where possible.  This helps to avoid
+/// issues with Euler flipping.
+template <typename T>
+std::vector<ModelParametersT<T>> transformPose(
+    const Character& character,
+    const std::vector<ModelParametersT<T>>& modelParameters,
+    const std::vector<TransformT<T>>& transforms,
+    bool ensureContinuousOutput = true);
+
+template <typename T>
+std::vector<ModelParametersT<T>> transformPose(
+    const Character& character,
+    const std::vector<ModelParametersT<T>>& modelParameters,
+    const TransformT<T>& transforms);
+
+/// Solver which can be reused to repeatedly apply the transformPose function listed above.
+template <typename T>
+class PoseTransformSolverT {
+ public:
+  explicit PoseTransformSolverT(const Character& character);
+  ~PoseTransformSolverT();
+
+  /// Applies a rigid transform to the given model parameters.
+  /// @param[in,out] modelParameters The model parameters to be transformed.
+  /// @param[in] transform The transform to apply.
+  /// @param[in] modelParametersPrev The previous pose, which is used to ensure continuity (the
+  /// solver will initialize using the previous frame's rigid parameters).
+  void transformPose(
+      ModelParametersT<T>& modelParameters,
+      const TransformT<T>& transform,
+      const ModelParametersT<T>& modelParametersPrev = ModelParametersT<T>{});
+
+ private:
+  Eigen::Index numModelParametersFull_;
+
+  momentum::Character characterSimplified_;
+  momentum::ParameterTransformT<T> parameterTransformSimplified_;
+  size_t rootJointSimplified_;
+  momentum::ParameterSet rigidParametersSimplified_;
+  std::vector<size_t> simplifiedParamToFullParamIdx_;
+
+  momentum::ModelParametersT<T> solvedParametersSimplified_;
+  momentum::SkeletonStateT<T> skelStateSimplified_;
+
+  std::shared_ptr<momentum::PositionErrorFunctionT<T>> positionError_;
+  std::shared_ptr<momentum::OrientationErrorFunctionT<T>> orientationError_;
+
+  std::unique_ptr<momentum::SkeletonSolverFunctionT<T>> solverFunction_;
+  std::unique_ptr<momentum::SolverT<T>> solver_;
+
+  momentum::Random<> rng_;
+};
+
+} // namespace momentum

include/momentum/character_solver/trust_region_qr.h

@@ -0,0 +1,78 @@
+/*
+ * 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/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_;
+
+  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,154 @@
+/*
+ * 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) final;
+
+  double getGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      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 {
+    return character_;
+  }
+
+ private:
+  double calculatePositionJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      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 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 VertexConstraintT<T>& constr,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  double calculateNormalGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      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_;
+
+  std::unique_ptr<MeshT<T>>
+      neutralMesh_; // Rest mesh without facial expression basis,
+                    // used to restore the neutral shape after facial expressions are applied.
+                    // Not used with there is a shape basis.
+  std::unique_ptr<MeshT<T>> restMesh_; // The rest positions of the mesh after shape basis
+                                       // (and potentially facial expression) has been applied
+  std::unique_ptr<MeshT<T>>
+      posedMesh_; // The posed mesh after the skeleton transforms have been applied.
+
+  const VertexConstraintType constraintType_;
+
+  uint32_t maxThreads_;
+
+  void updateMeshes(const ModelParametersT<T>& modelParameters, const SkeletonStateT<T>& state);
+};
+
+} // namespace momentum

include/momentum/character_solver/vertex_projection_error_function.h

@@ -0,0 +1,118 @@
+/*
+ * 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) final;
+
+  double getGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      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 {
+    return character_;
+  }
+
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+
+ private:
+  double calculateJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      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 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;
+
+  void updateMeshes(const ModelParametersT<T>& modelParameters, const SkeletonStateT<T>& state);
+
+  const Character& character_;
+
+  std::vector<VertexProjectionConstraintT<T>> constraints_;
+
+  std::unique_ptr<MeshT<T>>
+      neutralMesh_; // Rest mesh without facial expression basis,
+                    // used to restore the neutral shape after facial expressions are applied.
+                    // Not used with there is a shape basis.
+  std::unique_ptr<MeshT<T>> restMesh_; // The rest positions of the mesh after shape basis
+                                       // (and potentially facial expression) has been applied
+  std::unique_ptr<MeshT<T>>
+      posedMesh_; // The posed mesh after the skeleton transforms have been applied.
+
+  uint32_t maxThreads_;
+
+  T _nearClip = 1.0f;
+};
+
+} // namespace momentum