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

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

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