pymomentum-cpu 0.1.78.post11__cp313-cp313-win_amd64.whl

include/momentum/character_solver/simd_position_error_function.h

@@ -0,0 +1,165 @@
+/*
+ * 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/types.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/common/aligned.h>
+
+#include <array>
+#include <atomic>
+#include <cstdint>
+
+namespace momentum {
+
+struct SimdPositionConstraints final {
+ public:
+  explicit SimdPositionConstraints(const Skeleton* skel);
+
+  ~SimdPositionConstraints();
+
+  void clearConstraints();
+
+  void addConstraint(
+      size_t jointIndex,
+      const Vector3f& offset,
+      const Vector3f& target,
+      float targetWeight);
+
+  [[nodiscard]] VectorXi getNumConstraints() const;
+
+ public:
+  // max number of constraints per segment
+  static constexpr size_t kMaxConstraints = 4096;
+  static constexpr size_t kMaxJoints = 512;
+
+  // striped arrays for storing the offset
+  std::unique_ptr<float, AlignedDeleter> data;
+  float* offsetX;
+  float* offsetY;
+  float* offsetZ;
+  // striped arrays for storing the normal + target
+  float* targetX;
+  float* targetY;
+  float* targetZ;
+  // array for constraint weight
+  float* weights;
+
+  // store the number of constraints per joint
+  std::array<std::atomic<uint32_t>, kMaxJoints> constraintCount;
+  int numJoints;
+};
+
+/// A highly optimized error function for 3d position errors.
+///
+/// This class is best used when dealing with a large number of constraints. For fewer constraints,
+/// consider using the generic PositionErrorFunction instead.
+///
+/// @warning Due to the multi-threaded evaluation of the error/gradient, the functions are
+/// non-deterministic, which may cause numerical inconsistencies. As a result, slightly different
+/// results may be produced on multiple calls with the same data.
+class SimdPositionErrorFunction : public SkeletonErrorFunction {
+ public:
+  /// The dimensionality of each constraint.
+  static constexpr size_t kConstraintDim = 3;
+
+  /// @param maxThreads An optional parameter that specifies the maximum number of threads to be
+  /// used with dispenso::parallel_for. If this parameter is set to zero, the function will run in
+  /// serial mode, i.e., it will not use any additional threads. By default, the value is set to the
+  /// maximum allowable size of a uint32_t, which is also the default for dispenso.
+  explicit SimdPositionErrorFunction(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max());
+
+  /// @param maxThreads An optional parameter that specifies the maximum number of threads to be
+  /// used with dispenso::parallel_for. If this parameter is set to zero, the function will run in
+  /// serial mode, i.e., it will not use any additional threads. By default, the value is set to the
+  /// maximum allowable size of a uint32_t, which is also the default for dispenso.
+  explicit SimdPositionErrorFunction(
+      const Character& character,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max());
+
+  [[nodiscard]] double getError(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& /* unused */) override;
+
+  double getGradient(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& /* unused */,
+      Ref<VectorXf> gradient) override;
+
+  double getJacobian(
+      const ModelParameters& /* params */,
+      const SkeletonState& state,
+      const MeshState& /* unused */,
+      Ref<MatrixXf> jacobian,
+      Ref<VectorXf> residual,
+      int& usedRows) override;
+
+  [[nodiscard]] size_t getJacobianSize() const override;
+
+  void setConstraints(const SimdPositionConstraints* cstrs);
+
+ protected:
+  // weights for the error functions
+  static constexpr float kPositionWeight = 1e-4f;
+
+  mutable std::vector<size_t> jacobianOffset_;
+
+  // constraints to use
+  const SimdPositionConstraints* constraints_;
+
+  uint32_t maxThreads_;
+};
+
+#ifdef MOMENTUM_ENABLE_AVX
+
+// A version of SimdPositionErrorFunction where the Jacobian has been hand-unrolled using
+// AVX instructions.  On some platforms this performs better than the generic SIMD version
+// but it only works on Intel platforms that support AVX.
+class SimdPositionErrorFunctionAVX : public SimdPositionErrorFunction {
+ public:
+  explicit SimdPositionErrorFunctionAVX(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max());
+
+  explicit SimdPositionErrorFunctionAVX(
+      const Character& character,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max());
+
+  double getError(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& /* unused */) final;
+
+  double getGradient(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& /* unused */,
+      Ref<VectorXf> gradient) final;
+
+  double getJacobian(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& /* unused */,
+      Ref<MatrixXf> jacobian,
+      Ref<VectorXf> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+};
+
+#endif // MOMENTUM_ENABLE_AVX
+
+} // namespace momentum

include/momentum/character_solver/skeleton_error_function.h

@@ -0,0 +1,151 @@
+/*
+ * 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_solver/fwd.h>
+
+namespace momentum {
+
+template <typename T>
+class SkeletonErrorFunctionT {
+ public:
+  SkeletonErrorFunctionT(const Skeleton& skel, const ParameterTransform& pt)
+      : skeleton_(skel),
+        parameterTransform_(pt),
+        weight_(1.0),
+        activeJointParams_(pt.activeJointParams) {
+    enabledParameters_.flip(); // all parameters enabled by default
+  }
+  virtual ~SkeletonErrorFunctionT() = default;
+
+  [[nodiscard]] const Skeleton& getSkeleton() const {
+    return skeleton_;
+  }
+
+  [[nodiscard]] const ParameterTransform& getParameterTransform() const {
+    return parameterTransform_;
+  }
+
+  /// For mesh-based error functions, get access to the underlying character.  Error functions
+  /// that depend only on the skeleton can return nullptr.
+  [[nodiscard]] virtual const Character* getCharacter() const {
+    return nullptr;
+  }
+
+  void setWeight(T w) {
+    weight_ = w;
+  }
+
+  [[nodiscard]] T getWeight() const {
+    return weight_;
+  }
+
+  void setActiveJoints(const VectorX<bool>& aj) {
+    activeJointParams_ = aj;
+  }
+
+  void setEnabledParameters(const ParameterSet& ps) {
+    enabledParameters_ = ps;
+  }
+
+  virtual double getError(
+      const ModelParametersT<T>& /* params */,
+      const SkeletonStateT<T>& /* state */,
+      const MeshStateT<T>& /* meshState */) {
+    return 0.0f;
+  }
+
+  virtual double getGradient(
+      const ModelParametersT<T>& /* params */,
+      const SkeletonStateT<T>& /* state */,
+      const MeshStateT<T>& /* meshState */,
+      Eigen::Ref<Eigen::VectorX<T>> /* gradient */) {
+    return 0.0f;
+  }
+
+  virtual double getJacobian(
+      const ModelParametersT<T>& /* params */,
+      const SkeletonStateT<T>& /* state */,
+      const MeshStateT<T>& /* meshState */,
+      Eigen::Ref<Eigen::MatrixX<T>> /* jacobian */,
+      Eigen::Ref<Eigen::VectorX<T>> /* residual */,
+      int& usedRows) {
+    usedRows = 0;
+    return 0.0f;
+  }
+
+  virtual void getHessian(
+      const ModelParametersT<T>& /* params */,
+      const SkeletonStateT<T>& /* state */,
+      const MeshStateT<T>& /* meshState */,
+      Eigen::Ref<Eigen::MatrixX<T>> /* hessian */) {
+    throw;
+  }
+
+  virtual double getSolverDerivatives(
+      const ModelParametersT<T>& parameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      const size_t actualParameters,
+      Eigen::MatrixX<T>& hess,
+      Eigen::VectorX<T>& grad) {
+    const int paramSize = static_cast<int>(parameters.size());
+    double error = 0.0;
+
+    if (useHessian_) {
+      // get some ice cream, you deserve it
+      error = getGradient(parameters, state, meshState, grad);
+      getHessian(parameters, state, meshState, hess);
+    } else {
+      const int residualSize = static_cast<int>(getJacobianSize());
+
+      Eigen::MatrixX<T> jacobian = Eigen::MatrixX<T>::Zero(residualSize, paramSize);
+      Eigen::VectorX<T> residual = Eigen::VectorX<T>::Zero(residualSize);
+      int rows = 0;
+      error = getJacobian(parameters, state, meshState, jacobian, residual, rows);
+
+      // Update JtJ
+      if (rows > 0) {
+        // ! In truth, on the the "actualParameters" leftmost block will be used
+        // We take advantage of this here and skip the other computations
+        const auto JtBlock2 =
+            (jacobian.topLeftCorner(rows, actualParameters).transpose() * 2.0).eval();
+
+        // Update Hessian (Hessian += 2.0 * J^T * J) using selfadjointView with rankUpdate,
+        // replacing triangularView
+        hess.template selfadjointView<Eigen::Lower>().rankUpdate(JtBlock2);
+
+        // Update JtR
+        grad.noalias() += JtBlock2 * residual.head(rows);
+      }
+    }
+
+    return error;
+  }
+
+  [[nodiscard]] virtual size_t getJacobianSize() const {
+    return 0;
+  }
+
+  [[nodiscard]] virtual bool needsMesh() const {
+    return false;
+  }
+
+ protected:
+  const Skeleton& skeleton_;
+  const ParameterTransform& parameterTransform_;
+  T weight_;
+  VectorX<bool> activeJointParams_;
+  ParameterSet enabledParameters_; // set to zero by default constr
+  bool useHessian_{false}; // this can be true only if the getHessian is implemented
+};
+
+} // namespace momentum

include/momentum/character_solver/skeleton_solver_function.h

@@ -0,0 +1,94 @@
+/*
+ * 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/fwd.h>
+#include <momentum/character/parameter_transform.h>
+#include <momentum/character/types.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/solver/solver_function.h>
+#include <gsl/span>
+
+namespace momentum {
+
+template <typename T>
+class SkeletonSolverFunctionT : public SolverFunctionT<T> {
+ public:
+  SkeletonSolverFunctionT(
+      const Character& character,
+      const ParameterTransformT<T>& parameterTransform,
+      gsl::span<const std::shared_ptr<SkeletonErrorFunctionT<T>>> errorFunctions = {});
+  ~SkeletonSolverFunctionT() 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;
+
+  double getJtJR(
+      const Eigen::VectorX<T>& parameters,
+      Eigen::MatrixX<T>& jtj,
+      Eigen::VectorX<T>& jtr) override;
+
+  // overriding this to get a mix of JtJs and analytical Hessians from skeleton_ errorFunctions_
+  double getSolverDerivatives(
+      const Eigen::VectorX<T>& parameters,
+      Eigen::MatrixX<T>& hess,
+      Eigen::VectorX<T>& grad) override;
+
+  void updateParameters(Eigen::VectorX<T>& parameters, const Eigen::VectorX<T>& delta) final;
+  void setEnabledParameters(const ParameterSet& ps) final;
+
+  void addErrorFunction(std::shared_ptr<SkeletonErrorFunctionT<T>> solvable);
+  void clearErrorFunctions();
+
+  [[nodiscard]] const std::vector<std::shared_ptr<SkeletonErrorFunctionT<T>>>& getErrorFunctions()
+      const;
+
+  /// Provides access to the full character (replaces getSkeleton and getParameterTransform)
+  [[nodiscard]] const Character& getCharacter() const {
+    return character_;
+  }
+
+  /// Legacy compatibility (delegates to character_.skeleton)
+  [[nodiscard]] const Skeleton* getSkeleton() const {
+    return &character_.skeleton;
+  }
+
+  /// Provides access to the parameter transform
+  [[nodiscard]] const ParameterTransformT<T>* getParameterTransform() const {
+    return &parameterTransform_;
+  }
+
+  /// Updates mesh state if any error functions require it
+  void updateMeshState(const ModelParametersT<T>& parameters, const SkeletonStateT<T>& state);
+
+  /// Checks if any error functions require mesh state
+  [[nodiscard]] bool needsMeshState() const;
+
+ private:
+  const Character& character_;
+  const ParameterTransformT<T>& parameterTransform_;
+  std::unique_ptr<SkeletonStateT<T>> state_;
+  std::unique_ptr<MeshStateT<T>> meshState_;
+  bool needsMeshState_;
+  VectorX<bool> activeJointParams_;
+
+  Eigen::MatrixX<T> tJacobian_;
+  Eigen::VectorX<T> tResidual_;
+
+  std::vector<std::shared_ptr<SkeletonErrorFunctionT<T>>> errorFunctions_;
+};
+
+} // namespace momentum

include/momentum/character_solver/skinned_locator_error_function.h

@@ -0,0 +1,166 @@
+/*
+ * 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/skinned_locator.h>
+#include <momentum/character_solver/error_function_utils.h>
+#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 {
+
+/// Constraint structure for skinned locator error function.
+///
+/// Represents a constraint that pulls a skinned locator towards a target position
+/// in world space with a specified weight.
+template <typename T>
+struct SkinnedLocatorConstraintT {
+  int locatorIndex = -1;
+  T weight = 1;
+  Eigen::Vector3<T> targetPosition;
+
+  template <typename T2>
+  SkinnedLocatorConstraintT<T2> cast() const {
+    return {
+        this->locatorIndex,
+        static_cast<T2>(this->weight),
+        this->targetPosition.template cast<T2>()};
+  }
+};
+
+/// Error function for skinned locator constraints.
+///
+/// This error function computes the error between skinned locator positions and their
+/// target positions. Skinned locators are points that are deformed by the character's
+/// skeleton and can be used to constrain the character to specific world positions.
+/// The error is computed as the weighted squared distance between the current skinned
+/// locator position and the target position.
+template <typename T>
+class SkinnedLocatorErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  /// Constructs a skinned locator error function for the given character.
+  explicit SkinnedLocatorErrorFunctionT(const Character& character_in);
+  ~SkinnedLocatorErrorFunctionT() override;
+
+  SkinnedLocatorErrorFunctionT(const SkinnedLocatorErrorFunctionT<T>& other) = delete;
+  SkinnedLocatorErrorFunctionT<T>& operator=(const SkinnedLocatorErrorFunctionT<T>& other) = delete;
+  SkinnedLocatorErrorFunctionT(SkinnedLocatorErrorFunctionT<T>&& other) = delete;
+  SkinnedLocatorErrorFunctionT<T>& operator=(SkinnedLocatorErrorFunctionT<T>&& other) = delete;
+
+  /// Computes the total error for all skinned locator constraints.
+  ///
+  /// @param modelParameters The current model parameters
+  /// @param state The current skeleton state
+  /// @return The sum of weighted squared distances between locators and targets
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& /* meshState */) final;
+
+  /// Computes the gradient of the error function with respect to model parameters.
+  ///
+  /// @param modelParameters The current model parameters
+  /// @param state The current skeleton state
+  /// @param gradient Output gradient vector to be filled
+  /// @return The total error value
+  double getGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& /* meshState */,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  /// Computes the Jacobian matrix and residual vector for the error function.
+  ///
+  /// @param modelParameters The current model parameters
+  /// @param state The current skeleton state
+  /// @param jacobian Output Jacobian matrix to be filled
+  /// @param residual Output residual vector to be filled
+  /// @param usedRows Number of rows used in the Jacobian (output parameter)
+  /// @return The total error value
+  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;
+
+  /// Returns the number of rows needed for the Jacobian matrix.
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  /// Adds a new skinned locator constraint.
+  ///
+  /// @param locatorIndex Index of the skinned locator in the character
+  /// @param weight Weight for this constraint
+  /// @param targetPosition Target position in world space
+  void addConstraint(int locatorIndex, T weight, const Eigen::Vector3<T>& targetPosition);
+
+  /// Removes all constraints.
+  void clearConstraints();
+
+  /// Sets all constraints at once, replacing any existing constraints.
+  ///
+  /// @param constraints Vector of constraints to set
+  void setConstraints(const std::vector<SkinnedLocatorConstraintT<T>>& constraints);
+
+  /// Returns the current list of constraints.
+  [[nodiscard]] const std::vector<SkinnedLocatorConstraintT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+  /// Returns the character associated with this error function.
+  [[nodiscard]] const Character* getCharacter() const override {
+    return &character_;
+  }
+
+ private:
+  // Get the parameter index for a skinned locator, or -1 if not parameterized
+  [[nodiscard]] int getSkinnedLocatorParameterIndex(int locatorIndex) const {
+    if (locatorIndex <
+        static_cast<int>(character_.parameterTransform.skinnedLocatorParameters.size())) {
+      return character_.parameterTransform.skinnedLocatorParameters[locatorIndex];
+    }
+    return -1;
+  }
+
+  // Calculate the derivative of the world position with respect to the rest position
+  void calculateDWorldPos(
+      const SkeletonStateT<T>& state,
+      const SkinnedLocatorConstraintT<T>& constr,
+      const Eigen::Vector3<T>& d_restPos,
+      Eigen::Vector3<T>& d_worldPos) const;
+  double calculatePositionJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const SkinnedLocatorConstraintT<T>& constr,
+      Ref<Eigen::MatrixX<T>> jac,
+      Ref<Eigen::VectorX<T>> res) const;
+
+  double calculatePositionGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const SkinnedLocatorConstraintT<T>& constr,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  // Calculate the skinned locator position in world space
+  Eigen::Vector3<T> calculateSkinnedLocatorPosition(
+      const SkeletonStateT<T>& state,
+      int locatorIndex,
+      const Eigen::Vector3<T>& locatorRestPos) const;
+  Eigen::Vector3<T> getLocatorRestPosition(const ModelParametersT<T>& modelParams, int locatorIndex)
+      const;
+
+  const Character& character_;
+
+  std::vector<SkinnedLocatorConstraintT<T>> constraints_;
+};
+
+} // namespace momentum