pymomentum-cpu 0.1.77.post26__cp313-cp313-macosx_15_0_arm64.whl

include/momentum/character_solver/position_error_function.h

@@ -0,0 +1,75 @@
+/*
+ * 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/constraint_error_function.h>
+
+namespace momentum {
+
+/// 3D position constraint data
+template <typename T>
+struct PositionDataT : ConstraintData {
+  /// Positional offset in the parent joint space
+  Vector3<T> offset;
+  /// Target position
+  Vector3<T> target;
+
+  explicit PositionDataT(
+      const Vector3<T>& inOffset,
+      const Vector3<T>& inTarget,
+      size_t pIndex,
+      float w,
+      const std::string& n = "")
+      : ConstraintData(pIndex, w, n), offset(inOffset), target(inTarget) {}
+};
+
+/// The PositionErrorFunction computes the 3D positional errors from a list of Constraints.
+/// Each constraint specifies a locator on the skeleton (parent joint and offset), and its target 3D
+/// position.
+template <typename T>
+class PositionErrorFunctionT : public ConstraintErrorFunctionT<T, PositionDataT<T>> {
+ public:
+  /// Constructor
+  ///
+  /// @param[in] skel: character skeleton
+  /// @param[in] pt: character parameter transformation
+  /// @param[in] lossAlpha: alpha parameter for the loss function
+  /// @param[in] lossC: c parameter for the loss function
+  explicit PositionErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, PositionDataT<T>>(skel, pt, lossAlpha, lossC) {}
+
+  /// A convenience api where character contains info of the skeleton and parameter transform.
+  ///
+  /// @param[in] character: character definition
+  /// @param[in] lossAlpha: alpha parameter for the loss function
+  /// @param[in] lossC: c parameter for the loss function
+  explicit PositionErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : PositionErrorFunctionT(character.skeleton, character.parameterTransform, lossAlpha, lossC) {
+  }
+
+  /// Default constant weight in MarkerErrorFunction. This can be used for backwards compatibility
+  /// in setWeight().
+  static constexpr T kLegacyWeight = 1e-4f;
+
+ protected:
+  void evalFunction(
+      size_t constrIndex,
+      const JointStateT<T>& state,
+      Vector3<T>& f,
+      optional_ref<std::array<Vector3<T>, 1>> v = {},
+      optional_ref<std::array<Matrix3<T>, 1>> dfdv = {}) const final;
+};
+
+} // namespace momentum

include/momentum/character_solver/projection_error_function.h

@@ -0,0 +1,93 @@
+/*
+ * 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/locator.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+
+namespace momentum {
+
+template <typename T>
+struct ProjectionConstraintDataT {
+  // Project a 3d point (after skinning) to (rxz, ryz, z) where
+  // (rx,ry) is the residual.
+  //
+  // This encodes the target 2d point, the view and camera transforms, and the
+  // camera distortion (if any).
+  //
+  // This projection model was copied from the hand IK FullProjectionConstraint.
+  // Ideally we would refactor somehow to make them share more code but
+  // the kinematic models are so different this would take some work.
+  Eigen::Matrix<T, 3, 4> projection; // Projection matrix
+
+  size_t parent{}; // parent joint of the constraint
+  Eigen::Vector3<T> offset = Eigen::Vector3<T>::Zero(); // relative offset to the parent
+  T weight = 1; // constraint weight
+
+  Eigen::Vector2<T> target = Eigen::Vector2<T>::Zero();
+
+  static ProjectionConstraintDataT<T> createFromLocator(const momentum::Locator& locator);
+};
+
+template <typename T>
+class ProjectionErrorFunctionT : public momentum::SkeletonErrorFunctionT<T> {
+ public:
+  ProjectionErrorFunctionT(
+      const momentum::Skeleton& skel,
+      const momentum::ParameterTransform& pt,
+      T nearClip = T(1));
+
+  [[nodiscard]] double getError(
+      const momentum::ModelParametersT<T>& params,
+      const momentum::SkeletonStateT<T>& state,
+      const momentum::MeshStateT<T>& meshState) final;
+  double getGradient(
+      const momentum::ModelParametersT<T>& params,
+      const momentum::SkeletonStateT<T>& state,
+      const momentum::MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+  double getJacobian(
+      const momentum::ModelParametersT<T>& params,
+      const momentum::SkeletonStateT<T>& state,
+      const momentum::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(const ProjectionConstraintDataT<T>& data);
+
+  void clearConstraints() {
+    constraints_.clear();
+  }
+  [[nodiscard]] bool empty() const {
+    return constraints_.empty();
+  }
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+  void setConstraints(std::vector<ProjectionConstraintDataT<T>> constraints) {
+    constraints_ = std::move(constraints);
+  }
+
+  [[nodiscard]] const std::vector<ProjectionConstraintDataT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+ protected:
+  std::vector<ProjectionConstraintDataT<T>> constraints_;
+
+  // Projection error is roughly in radians, so a value near 1 is reasonable here:
+  static constexpr T kProjectionWeight = 1.0f;
+
+  // Ignore projection constraints involving joints closer than this distance.
+  // Prevents divide-by-zero in the projection matrix and bad behavior close to
+  // the camera.
+  T _nearClip = 1.0f;
+};
+
+} // namespace momentum

include/momentum/character_solver/simd_collision_error_function.h

@@ -0,0 +1,99 @@
+/*
+ * 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/collision_geometry.h>
+#include <momentum/character/collision_geometry_state.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/common/aligned.h>
+#include <momentum/simd/simd.h>
+
+#include <axel/Bvh.h>
+
+#include <unordered_set>
+#include <vector>
+
+namespace momentum {
+
+/// Represents an error function that penalizes self-intersections between collision geometries of
+/// the character.
+///
+/// The function pre-filters collision geometries to ignore those attached to the same joint, those
+/// intersecting in the rest pose, or those very close in the rest pose.
+template <typename T>
+class SimdCollisionErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  explicit SimdCollisionErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const CollisionGeometry& collisionGeometry);
+
+  explicit SimdCollisionErrorFunctionT(const Character& character);
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& /* meshState */) final;
+
+  double getGradient(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& /* meshState */,
+      Ref<VectorX<T>> gradient) override;
+
+  double getJacobian(
+      const ModelParametersT<T>& /*unused*/,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& /* meshState */,
+      Ref<MatrixX<T>> jacobian,
+      Ref<VectorX<T>> residual,
+      int& usedRows) override;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+ protected:
+  void updateCollisionPairs();
+
+  // Update collisionState_, bvh_, and collisionPairs_ given the new skeleton state.
+  void computeBroadPhase(const SkeletonStateT<T>& state);
+
+  using PairId = std::pair<size_t, size_t>;
+
+  struct PairHash {
+    std::size_t operator()(const std::pair<size_t, size_t>& p) const {
+      // Use the smaller value as the high-order bits and the larger value as the low-order bits
+      return (static_cast<uint64_t>(std::min(p.first, p.second)) << (sizeof(size_t) * 8 / 2)) |
+          std::max(p.first, p.second);
+    }
+  };
+
+  struct PairEqual {
+    bool operator()(const PairId& a, const PairId& b) const {
+      return (a.first == b.first && a.second == b.second) ||
+          (a.first == b.second && a.second == b.first);
+    }
+  };
+
+  const CollisionGeometry collisionGeometry_;
+
+  std::unordered_set<PairId, PairHash, PairEqual> excludingPairIds_;
+
+  std::vector<std::vector<int, AlignedAllocator<int, momentum::kSimdAlignment>>> collisionPairs_;
+
+  size_t jacobianSize_ = 0;
+
+  CollisionGeometryStateT<T> collisionState_;
+
+  // weights for the error functions
+  static constexpr T kCollisionWeight = 5e-3f;
+
+  axel::Bvh<T> bvh_;
+};
+
+} // namespace momentum

include/momentum/character_solver/simd_normal_error_function.h

@@ -0,0 +1,157 @@
+/*
+ * 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_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 SimdNormalConstraints {
+ public:
+  explicit SimdNormalConstraints(const Skeleton* skel);
+  ~SimdNormalConstraints();
+
+  void clearConstraints();
+
+  bool addConstraint(
+      size_t jointIndex,
+      const Vector3f& offset,
+      const Vector3f& normal,
+      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* normalX;
+  float* normalY;
+  float* normalZ;
+  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 "point-to-plane" (signed) distance errors. See
+/// NormalErrorFunction for detailed explanation.
+///
+/// This function should primarily be used when dealing with a large number of constraints. For
+/// smaller numbers of constraints, consider using the generic NormalErrorFunction.
+///
+/// @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 occur on multiple calls with the same data.
+class SimdNormalErrorFunction : public SkeletonErrorFunction {
+ public:
+  /// @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 SimdNormalErrorFunction(
+      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 SimdNormalErrorFunction(
+      const Character& character,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max());
+
+  [[nodiscard]] double getError(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& meshState) final;
+
+  double getGradient(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& meshState,
+      Ref<VectorXf> gradient) final;
+
+  double getJacobian(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& meshState,
+      Ref<MatrixXf> jacobian,
+      Ref<VectorXf> residual,
+      int& usedRows) override;
+
+  [[nodiscard]] size_t getJacobianSize() const override;
+
+  void setConstraints(const SimdNormalConstraints* cstrs) {
+    constraints_ = cstrs;
+  }
+
+ protected:
+  // weights for the error functions
+  static constexpr float kPlaneWeight = 1e-4f;
+
+  mutable std::vector<size_t> jacobianOffset_;
+
+  // constraints to use
+  const SimdNormalConstraints* constraints_;
+
+  uint32_t maxThreads_;
+};
+
+#ifdef MOMENTUM_ENABLE_AVX
+
+// A version of SimdNormalErrorFunction 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 SimdNormalErrorFunctionAVX : public SimdNormalErrorFunction {
+ public:
+  explicit SimdNormalErrorFunctionAVX(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max())
+      : SimdNormalErrorFunction(skel, pt, maxThreads) {}
+  explicit SimdNormalErrorFunctionAVX(
+      const Character& character,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max())
+      : SimdNormalErrorFunction(character, maxThreads) {}
+
+  double getJacobian(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& meshState,
+      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/simd_plane_error_function.h

@@ -0,0 +1,164 @@
+/*
+ * 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_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/common/aligned.h>
+
+#include <array>
+#include <atomic>
+#include <cstdint>
+#include <vector>
+
+namespace momentum {
+
+struct SimdPlaneConstraints final {
+ public:
+  explicit SimdPlaneConstraints(const Skeleton* skel);
+
+  ~SimdPlaneConstraints();
+
+  void clearConstraints();
+
+  void addConstraint(
+      size_t jointIndex,
+      const Vector3f& offset,
+      const Vector3f& targetNormal,
+      float targetOffset,
+      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* normalX;
+  float* normalY;
+  float* normalZ;
+  float* targets;
+  // 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 "point-on-plane" or "point-in-half-plane" errors.
+///
+/// This function is recommended for use only when dealing with a large number of constraints. For a
+/// smaller number of constraints, consider using the generic PlaneErrorFunction.
+///
+/// @warning Due to the multi-threaded evaluation of the error/gradient, the functions are
+/// non-deterministic. This might lead to numerical inconsistencies, resulting in slightly different
+/// outcomes on multiple calls with the same data.
+class SimdPlaneErrorFunction : public SkeletonErrorFunction {
+ public:
+  /// @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 SimdPlaneErrorFunction(
+      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 SimdPlaneErrorFunction(
+      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& /* meshState */,
+      Ref<VectorXf> gradient) override;
+
+  double getJacobian(
+      const ModelParameters& params,
+      const SkeletonState& state,
+      const MeshState& /* meshState */,
+      Ref<MatrixXf> jacobian,
+      Ref<VectorXf> residual,
+      int& usedRows) override;
+
+  [[nodiscard]] size_t getJacobianSize() const override;
+
+  void setConstraints(const SimdPlaneConstraints* cstrs);
+
+ protected:
+  // weights for the error functions
+  static constexpr float kPlaneWeight = 1e-4f;
+
+  mutable std::vector<size_t> jacobianOffset_;
+
+  // constraints to use
+  const SimdPlaneConstraints* constraints_;
+
+  uint32_t maxThreads_;
+};
+
+#ifdef MOMENTUM_ENABLE_AVX
+
+// A version of SimdPlaneErrorFunction 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 SimdPlaneErrorFunctionAVX : public SimdPlaneErrorFunction {
+ public:
+  explicit SimdPlaneErrorFunctionAVX(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max());
+
+  explicit SimdPlaneErrorFunctionAVX(
+      const Character& character,
+      uint32_t maxThreads = std::numeric_limits<uint32_t>::max());
+
+  [[nodiscard]] 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