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

include/momentum/character_sequence_solver/multipose_solver_function.h

@@ -0,0 +1,82 @@
+/*
+ * 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/parameter_transform.h>
+#include <momentum/character/types.h>
+#include <momentum/character_sequence_solver/fwd.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/solver/solver_function.h>
+
+namespace momentum {
+
+template <typename T>
+class MultiposeSolverFunctionT : public SolverFunctionT<T> {
+ public:
+  MultiposeSolverFunctionT(
+      const Skeleton* skel,
+      const ParameterTransformT<T>* parameterTransform,
+      std::span<const int> universal,
+      size_t frames);
+  ~MultiposeSolverFunctionT() override;
+
+  double getError(const Eigen::VectorX<T>& parameters) final;
+
+  double getGradient(const Eigen::VectorX<T>& parameters, Eigen::VectorX<T>& gradient) final;
+
+  double getJacobian(
+      const Eigen::VectorX<T>& parameters,
+      Eigen::MatrixX<T>& jacobian,
+      Eigen::VectorX<T>& residual,
+      size_t& actualRows) final;
+
+  void updateParameters(Eigen::VectorX<T>& parameters, const Eigen::VectorX<T>& gradient) final;
+  void setEnabledParameters(const ParameterSet& parameterSet) final;
+
+  void addErrorFunction(size_t frame, SkeletonErrorFunctionT<T>* errorFunction);
+
+  [[nodiscard]] size_t getNumFrames() const {
+    return states_.size();
+  }
+
+  [[nodiscard]] const ModelParametersT<T>& getFrameParameters(const size_t frame) const {
+    return frameParameters_[frame];
+  }
+
+  void setFrameParameters(size_t frame, const ModelParametersT<T>& parameters);
+  [[nodiscard]] Eigen::VectorX<T> getUniversalParameters() const;
+  [[nodiscard]] Eigen::VectorX<T> getJoinedParameterVector() const;
+  void setJoinedParameterVector(const Eigen::VectorX<T>& joinedParameters);
+
+  [[nodiscard]] ParameterSet getUniversalParameterSet() const;
+  [[nodiscard]] ParameterSet getUniversalLocatorParameterSet() const;
+
+ private:
+  void setFrameParametersFromJoinedParameterVector(const Eigen::VectorX<T>& parameters);
+
+ private:
+  const Skeleton* skeleton_;
+  const ParameterTransformT<T>* parameterTransform_;
+  std::vector<SkeletonStateT<T>> states_;
+  std::vector<MeshStateT<T>> meshStates_;
+  VectorX<bool> activeJointParams_;
+
+  std::vector<ModelParametersT<T>> frameParameters_;
+  Eigen::VectorX<T> universal_;
+  std::vector<size_t> parameterIndexMap_;
+
+  std::vector<size_t> genericParameters_;
+  std::vector<size_t> universalParameters_;
+
+  std::vector<std::vector<SkeletonErrorFunctionT<T>*>> errorFunctions_;
+
+  friend class MultiposeSolverT<T>;
+};
+
+} // namespace momentum

include/momentum/character_sequence_solver/sequence_error_function.h

@@ -0,0 +1,104 @@
+/*
+ * 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/parameter_transform.h>
+#include <momentum/character/types.h>
+#include <momentum/character_sequence_solver/fwd.h>
+
+namespace momentum {
+
+template <typename T>
+class SequenceErrorFunctionT {
+ public:
+  SequenceErrorFunctionT(const Skeleton& skel, const ParameterTransform& pt)
+      : skeleton_(skel), parameterTransform_(pt), activeJointParams_(pt.activeJointParams) {
+    enabledParameters_.flip(); // all parameters enabled by default
+  }
+  virtual ~SequenceErrorFunctionT() = default;
+
+  // Number of contiguous frames that this error function affects.
+  [[nodiscard]] virtual size_t numFrames() const = 0;
+
+  void setWeight(T w) {
+    weight_ = w;
+  }
+  [[nodiscard]] float getWeight() const {
+    return weight_;
+  }
+
+  void setActiveJoints(const VectorX<bool>& aj) {
+    activeJointParams_ = aj;
+  }
+
+  void setEnabledParameters(const ParameterSet& ps) {
+    enabledParameters_ = ps;
+  }
+
+  [[nodiscard]] virtual bool needsMesh() const {
+    return false;
+  }
+
+  [[nodiscard]] virtual double getError(
+      std::span<const ModelParametersT<T>> /* modelParameters */,
+      std::span<const SkeletonStateT<T>> /* skelStates */,
+      std::span<const MeshStateT<T>> /* meshStates */) const {
+    return 0.0f;
+  }
+
+  // Get the gradient of the error.
+  // modelParameters: numFrames() array of parameter vectors
+  // skelStates: [numFrames()] array of skeleton states
+  // meshStates: [numFrames()] array of mesh states
+  // gradient: [numFrames() * parameterTransform] gradient vector
+  virtual double getGradient(
+      std::span<const ModelParametersT<T>> /* modelParameters */,
+      std::span<const SkeletonStateT<T>> /* skelStates */,
+      std::span<const MeshStateT<T>> /* meshStates */,
+      Eigen::Ref<Eigen::VectorX<T>> /* gradient */) const {
+    return 0.0f;
+  }
+
+  // modelParameters: [numFrames() * parameterTransform] parameter vector
+  // skelStates: [numFrames()] array of skeleton states
+  // meshStates: [numFrames()] array of mesh states
+  // jacobian: [getJacobianSize()] x [numFrames() * parameterTransform] Jacobian matrix
+  // residual: [getJacobianSize()] residual vector.
+  virtual double getJacobian(
+      std::span<const ModelParametersT<T>> /* modelParameters */,
+      std::span<const SkeletonStateT<T>> /* skelStates */,
+      std::span<const MeshStateT<T>> /* meshStates */,
+      Eigen::Ref<Eigen::MatrixX<T>> /* jacobian */,
+      Eigen::Ref<Eigen::VectorX<T>> /* residual */,
+      int& usedRows) const {
+    usedRows = 0;
+    return 0.0f;
+  }
+
+  [[nodiscard]] virtual size_t getJacobianSize() const {
+    return 0;
+  }
+
+  [[nodiscard]] const Skeleton& getSkeleton() const {
+    return this->skeleton_;
+  }
+
+  [[nodiscard]] const ParameterTransform& getParameterTransform() const {
+    return this->parameterTransform_;
+  }
+
+ protected:
+  const Skeleton& skeleton_;
+  const ParameterTransform& parameterTransform_;
+  T weight_ = 1.0f;
+  VectorX<bool> activeJointParams_;
+  ParameterSet enabledParameters_; // set to zero by default constr
+};
+
+} // namespace momentum

include/momentum/character_sequence_solver/sequence_solver.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/fwd.h>
+#include <momentum/character_sequence_solver/fwd.h>
+#include <momentum/common/fwd.h>
+#include <momentum/math/fwd.h>
+#include <momentum/math/online_householder_qr.h>
+#include <momentum/solver/solver.h>
+
+#include <functional>
+#include <span>
+
+namespace momentum {
+
+/// Sequence solver specific options
+struct SequenceSolverOptions : SolverOptions {
+  /// Regularization parameter for QR decomposition.
+  float regularization = 0.05f;
+
+  /// Flag to enable line search during optimization.
+  bool doLineSearch = false;
+
+  /// Flag to enable multithreading for the Sequence solver.
+  bool multithreaded = false;
+
+  /// Flag to enable a progress bar during optimization.
+  bool progressBar = false;
+
+  SequenceSolverOptions() = default;
+
+  /* implicit */ SequenceSolverOptions(const SolverOptions& baseOptions)
+      : SolverOptions(baseOptions) {
+    // Empty
+  }
+};
+
+template <typename T>
+class SequenceSolverT : public SolverT<T> {
+ public:
+  SequenceSolverT(const SolverOptions& options, SequenceSolverFunctionT<T>* function);
+
+  [[nodiscard]] std::string_view getName() const override;
+
+  void setOptions(const SolverOptions& options) final;
+
+  // "numParameters" has a different meaning here than in Solver.
+  // For Solver class, numParameters = ParameterTransform.numAllModelParameters() * numFrames; but
+  // it means #variables here. To hack around this discrepancy, we update numParameters when a new
+  // set of variables are defined.
+  // XXX To be fixed as part of T118191244
+  void setEnabledParameters(const ParameterSet& parameters) final {
+    SolverT<T>::setEnabledParameters(parameters);
+    this->numParameters_ = this->solverFunction_->getNumParameters();
+  }
+
+  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:
+  static double processPerFrameErrors_parallel(
+      SequenceSolverFunctionT<T>* fn,
+      OnlineBandedHouseholderQR<T>& qrSolver,
+      ProgressBar* progress);
+  static double processPerFrameErrors_serial(
+      SequenceSolverFunctionT<T>* fn,
+      OnlineBandedHouseholderQR<T>& qrSolver,
+      ProgressBar* progress);
+
+  static double processSequenceErrors_serial(
+      SequenceSolverFunctionT<T>* fn,
+      OnlineBandedHouseholderQR<T>& qrSolver,
+      ProgressBar* progress);
+
+  struct JacobianResidual {
+    size_t frameIndex = SIZE_MAX;
+    Eigen::MatrixX<T> jacobian;
+    Eigen::VectorX<T> residual;
+    double error;
+    size_t nFunctions;
+    Eigen::Index usedRows = 0;
+
+    bool operator<(const JacobianResidual& rhs) const {
+      return frameIndex < rhs.frameIndex;
+    }
+
+    bool operator>(const JacobianResidual& rhs) const {
+      return frameIndex > rhs.frameIndex;
+    }
+  };
+
+  // Here, the processJac function computes the Jacobian/residual, applies it to the banded
+  // part of the matrix, and then returns the universal part.
+  static double processErrorFunctions_parallel(
+      const std::function<JacobianResidual(
+          size_t,
+          SequenceSolverFunctionT<T>*,
+          OnlineBandedHouseholderQR<T>&,
+          SkeletonStateT<T> skelState,
+          MeshStateT<T>& meshState)>& processJac,
+      SequenceSolverFunctionT<T>* fn,
+      OnlineBandedHouseholderQR<T>& qrSolver,
+      ProgressBar* progress);
+
+  // Returns the [Jacobian, residual, error] for all the error functions applying to a single frame:
+  static JacobianResidual computePerFrameJacobian(
+      SequenceSolverFunctionT<T>* fn,
+      size_t iFrame,
+      SkeletonStateT<T>& skelState,
+      MeshStateT<T>& meshState);
+
+  // Returns the [Jacobian, residual, error] for all the sequence error functions starting from a
+  // single frame:
+  static JacobianResidual computeSequenceJacobian(
+      SequenceSolverFunctionT<T>* fn,
+      size_t iFrame,
+      size_t bandwidth,
+      std::span<const SkeletonStateT<T>> skelStates,
+      std::span<const MeshStateT<T>> meshStates);
+
+  static std::vector<Eigen::Index> buildSequenceColumnIndices(
+      const SequenceSolverFunctionT<T>* fn,
+      size_t bandwidth);
+
+  float regularization_ = 0.05f;
+  bool doLineSearch_ = false;
+  bool multithreaded_ = false;
+  bool progressBar_ = false;
+
+  // bandwidth in frames:
+  size_t bandwidth_ = 0;
+};
+
+} // namespace momentum

include/momentum/character_sequence_solver/sequence_solver_function.h

@@ -0,0 +1,158 @@
+/*
+ * 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/parameter_transform.h>
+#include <momentum/character/types.h>
+#include <momentum/character_sequence_solver/fwd.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/solver/solver_function.h>
+
+#include <atomic>
+
+namespace momentum {
+
+template <typename T>
+struct PerFrameStateT {
+  SkeletonStateT<T> skeletonState;
+  MeshStateT<T> meshState;
+};
+
+inline constexpr size_t kAllFrames = SIZE_MAX;
+
+template <typename T>
+class SequenceSolverFunctionT : public SolverFunctionT<T> {
+ public:
+  SequenceSolverFunctionT(
+      const Character& character,
+      const ParameterTransformT<T>& parameterTransform,
+      const ParameterSet& universal,
+      size_t nFrames);
+  ~SequenceSolverFunctionT() override;
+
+  double getError(const Eigen::VectorX<T>& parameters) final;
+
+  double getGradient(const Eigen::VectorX<T>& parameters, Eigen::VectorX<T>& gradient) final;
+
+  double getJacobian(
+      const Eigen::VectorX<T>& parameters,
+      Eigen::MatrixX<T>& jacobian,
+      Eigen::VectorX<T>& residual,
+      size_t& actualRows) final;
+
+  void updateParameters(Eigen::VectorX<T>& parameters, const Eigen::VectorX<T>& gradient) final;
+  void setEnabledParameters(const ParameterSet& parameterSet) final;
+
+  [[nodiscard]] const ParameterSet& getUniversalParameterSet() const {
+    return universalParameters_;
+  }
+
+  /// Returns whether the per-frame error functions need the mesh.
+  [[nodiscard]] bool needsMeshPerFrame() const {
+    return needsMeshPerFrame_;
+  }
+
+  /// Returns whether the sequence error functions need the mesh.
+  [[nodiscard]] bool needsMeshSequence() const {
+    return needsMeshSequence_;
+  }
+
+  // Passing in the special frame index kAllFrames will add the error function to every frame; this
+  // is convenient for e.g. limit errors but requires that the error function be stateless. Note:
+  // you are allowed to call this in a multithreaded context but you must ensure the frame indices
+  // are different between the different threads.
+  void addErrorFunction(size_t frame, std::shared_ptr<SkeletonErrorFunctionT<T>> errorFunction);
+  void addSequenceErrorFunction(
+      size_t startFrame,
+      std::shared_ptr<SequenceErrorFunctionT<T>> errorFunction);
+
+  [[nodiscard]] size_t getNumFrames() const {
+    return frameParameters_.size();
+  }
+
+  [[nodiscard]] const ModelParametersT<T>& getFrameParameters(size_t frame) const {
+    return frameParameters_[frame];
+  }
+
+  void setFrameParameters(size_t frame, const ModelParametersT<T>& parameters);
+  [[nodiscard]] ModelParametersT<T> getUniversalParameters() const;
+  [[nodiscard]] Eigen::VectorX<T> getJoinedParameterVector() const;
+
+  /// Returns a joined parameter vector from a span of frame parameters.
+  ///
+  /// @param frameParameters A span of frame parameters, where each element is a ModelParametersT<T>
+  /// object representing the parameters for a single frame.
+  ///
+  /// @return An Eigen::VectorX<T> object containing the joined parameter vector of all frames.
+  [[nodiscard]] Eigen::VectorX<T> getJoinedParameterVectorFromFrameParameters(
+      std::span<const ModelParametersT<T>> frameParameters) const;
+
+  void setJoinedParameterVector(const Eigen::VectorX<T>& joinedParameters);
+
+  [[nodiscard]] const Skeleton* getSkeleton() const;
+  [[nodiscard]] const Character& getCharacter() const;
+
+  [[nodiscard]] const ParameterTransformT<T>* getParameterTransform() const {
+    return &parameterTransform_;
+  }
+
+  [[nodiscard]] const auto& getErrorFunctions(size_t iFrame) const {
+    return perFrameErrorFunctions_.at(iFrame);
+  }
+
+  [[nodiscard]] const auto& getSequenceErrorFunctions(size_t iFrame) const {
+    return sequenceErrorFunctions_.at(iFrame);
+  }
+
+ private:
+  void setFrameParametersFromJoinedParameterVector(const Eigen::VectorX<T>& parameters);
+
+  /// Helper method to compute per-frame jacobians
+  double
+  getPerFrameJacobian(Eigen::MatrixX<T>& jacobian, Eigen::VectorX<T>& residual, size_t& position);
+
+  /// Helper method to compute sequence error functions jacobians
+  double getSequenceErrorFunctionsJacobian(
+      Eigen::MatrixX<T>& jacobian,
+      Eigen::VectorX<T>& residual,
+      size_t& position);
+
+ private:
+  const Character& character_;
+  const ParameterTransformT<T>& parameterTransform_;
+  VectorX<bool> activeJointParams_;
+
+  std::vector<ModelParametersT<T>> frameParameters_;
+  size_t nFrames_;
+
+  void updateParameterSets(const ParameterSet& activeParams);
+
+  // Indices of parameters that are active and solved per-frame:
+  std::vector<Eigen::Index> perFrameParameterIndices_;
+
+  // Indices of parameters that are active and solved universally:
+  std::vector<Eigen::Index> universalParameterIndices_;
+
+  const ParameterSet universalParameters_;
+
+  std::vector<std::vector<std::shared_ptr<SkeletonErrorFunctionT<T>>>> perFrameErrorFunctions_;
+  std::vector<std::vector<std::shared_ptr<SequenceErrorFunctionT<T>>>> sequenceErrorFunctions_;
+
+  std::atomic<size_t> numTotalPerFrameErrorFunctions_ = 0;
+  std::atomic<size_t> numTotalSequenceErrorFunctions_ = 0;
+
+  // Whether per-frame error functions need the mesh
+  std::atomic<bool> needsMeshPerFrame_{false};
+  // Whether sequence error functions need the mesh
+  std::atomic<bool> needsMeshSequence_{false};
+
+  friend class SequenceSolverT<T>;
+};
+
+} // namespace momentum

include/momentum/character_sequence_solver/state_sequence_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/skeleton_state.h>
+#include <momentum/character_sequence_solver/fwd.h>
+#include <momentum/character_sequence_solver/sequence_error_function.h>
+#include <momentum/character_solver/state_error_function.h>
+
+namespace momentum {
+
+/// Error function that penalizes the difference between the current skeleton state and the next
+/// skeleton state. Note that by default this penalizes any difference between adjacent skeleton
+/// states. If we specify a target state, this will be applied as an offset to the current state,
+/// that is, the error will penalize the difference between offset * currentState and nextState.
+template <typename T>
+class StateSequenceErrorFunctionT : public SequenceErrorFunctionT<T> {
+ public:
+  StateSequenceErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      RotationErrorType rotationErrorType = RotationErrorType::RotationMatrixDifference);
+  explicit StateSequenceErrorFunctionT(
+      const Character& character,
+      RotationErrorType rotationErrorType = RotationErrorType::RotationMatrixDifference);
+
+  [[nodiscard]] size_t numFrames() const final {
+    return 2;
+  }
+
+  double getError(
+      std::span<const ModelParametersT<T>> modelParameters,
+      std::span<const SkeletonStateT<T>> skelStates,
+      std::span<const MeshStateT<T>> /* meshStates */) const final;
+  double getGradient(
+      std::span<const ModelParametersT<T>> modelParameters,
+      std::span<const SkeletonStateT<T>> skelStates,
+      std::span<const MeshStateT<T>> /* meshStates */,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const final;
+
+  // modelParameters: [numFrames() * parameterTransform] parameter vector
+  // skelStates: [numFrames()] array of skeleton states
+  // jacobian: [getJacobianSize()] x [numFrames() * parameterTransform] Jacobian matrix
+  // residual: [getJacobianSize()] residual vector.
+  double getJacobian(
+      std::span<const ModelParametersT<T>> modelParameters,
+      std::span<const SkeletonStateT<T>> skelStates,
+      std::span<const MeshStateT<T>> /* meshStates */,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) const final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  /// Set the target weights for each joint's position and rotation.
+  /// @param posWeight Per-joint position weights.
+  /// @param rotWeight Per-joint rotation 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;
+  }
+
+  void reset();
+
+  [[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_;
+  }
+
+  /// Set the target state for the skeleton.
+  ///
+  /// This function allows specifying a target state that the skeleton should aim to achieve.
+  /// The target state is represented as a list of transforms, one for each joint.
+  /// The error will penalize the difference between offset * currentState and nextState.
+  /// @param target A list of transforms representing the target state for each joint.
+  void setTargetState(TransformListT<T> target);
+
+ private:
+  Eigen::VectorX<T> targetPositionWeights_;
+  Eigen::VectorX<T> targetRotationWeights_;
+
+  TransformListT<T> targetState_;
+
+  T posWgt_;
+  T rotWgt_;
+
+  const RotationErrorType rotationErrorType_;
+
+ public:
+  // weights for the error functions
+  static constexpr T kPositionWeight = 1e-3f;
+  static constexpr T kOrientationWeight = 1e+0f;
+};
+
+} // namespace momentum