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

include/momentum/character_solver/error_function_utils.h

@@ -0,0 +1,60 @@
+/*
+ * 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>
+
+namespace momentum {
+
+template <typename T>
+void gradient_jointParams_to_modelParams(
+    const T& grad_jointParams,
+    const Eigen::Index iJointParam,
+    const ParameterTransform& parameterTransform,
+    Eigen::Ref<Eigen::VectorX<T>> gradient) {
+  // explicitly multiply with the parameter transform to generate parameter space gradients
+  for (auto index = parameterTransform.transform.outerIndexPtr()[iJointParam];
+       index < parameterTransform.transform.outerIndexPtr()[iJointParam + 1];
+       ++index) {
+    gradient[parameterTransform.transform.innerIndexPtr()[index]] +=
+        grad_jointParams * parameterTransform.transform.valuePtr()[index];
+  }
+}
+
+template <typename T>
+void jacobian_jointParams_to_modelParams(
+    const Eigen::Ref<const Eigen::VectorX<T>>& jacobian_jointParams,
+    const Eigen::Index iJointParam,
+    const ParameterTransform& parameterTransform,
+    Eigen::Ref<Eigen::MatrixX<T>> jacobian) {
+  // explicitly multiply with the parameter transform to generate parameter space gradients
+  for (auto index = parameterTransform.transform.outerIndexPtr()[iJointParam];
+       index < parameterTransform.transform.outerIndexPtr()[iJointParam + 1];
+       ++index) {
+    jacobian.col(parameterTransform.transform.innerIndexPtr()[index]) +=
+        jacobian_jointParams * parameterTransform.transform.valuePtr()[index];
+  }
+}
+
+template <typename T>
+void jacobian_jointParams_to_modelParams(
+    const T& jacobian_jointParams,
+    const Eigen::Index iJointParam,
+    const ParameterTransform& parameterTransform,
+    Eigen::Ref<Eigen::MatrixX<T>> jacobian) {
+  // explicitly multiply with the parameter transform to generate parameter space gradients
+  for (auto index = parameterTransform.transform.outerIndexPtr()[iJointParam];
+       index < parameterTransform.transform.outerIndexPtr()[iJointParam + 1];
+       ++index) {
+    jacobian(0, parameterTransform.transform.innerIndexPtr()[index]) +=
+        jacobian_jointParams * parameterTransform.transform.valuePtr()[index];
+  }
+}
+
+} // namespace momentum

include/momentum/character_solver/fixed_axis_error_function.h

@@ -0,0 +1,139 @@
+/*
+ * 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 {
+
+/// Axis constraint data
+template <typename T>
+struct FixedAxisDataT : ConstraintData {
+  /// axis in the parent's coordinate frame
+  Vector3<T> localAxis;
+  /// target axis in the global frame
+  Vector3<T> globalAxis;
+
+  explicit FixedAxisDataT(
+      const Vector3<T>& local,
+      const Vector3<T>& global,
+      size_t pIndex,
+      float w,
+      const std::string& n = "")
+      : ConstraintData(pIndex, w, n),
+        localAxis(local.normalized()),
+        globalAxis(global.normalized()) {}
+};
+
+// The fixed angle error minimizes the difference between a local axis and a target direction in the
+// global frame. These error functions differ in how they calculate the difference between two axis.
+
+/// Diff: compute the element wise differences between the two axis
+template <typename T>
+class FixedAxisDiffErrorFunctionT : public ConstraintErrorFunctionT<T, FixedAxisDataT<T>, 3, 1, 0> {
+ public:
+  explicit FixedAxisDiffErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, FixedAxisDataT<T>, 3, 1, 0>(skel, pt, lossAlpha, lossC) {}
+
+  explicit FixedAxisDiffErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : FixedAxisDiffErrorFunctionT<T>(
+            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-1f;
+
+ 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;
+};
+
+/// Cos: compute the dot product between the two axis - cosine of the angle between them
+template <typename T>
+class FixedAxisCosErrorFunctionT : public ConstraintErrorFunctionT<T, FixedAxisDataT<T>, 1, 1, 0> {
+ public:
+  explicit FixedAxisCosErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, FixedAxisDataT<T>, 1, 1, 0>(skel, pt, lossAlpha, lossC) {}
+
+  explicit FixedAxisCosErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : FixedAxisCosErrorFunctionT<T>(
+            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-1f;
+
+ protected:
+  void evalFunction(
+      size_t constrIndex,
+      const JointStateT<T>& state,
+      Vector<T, 1>& f,
+      optional_ref<std::array<Vector3<T>, 1>> v = {},
+      optional_ref<std::array<Eigen::Matrix<T, 1, 3>, 1>> dfdv = {}) const final;
+};
+
+/// Angle: compute the angle between the two axis using acosine of their dot product.
+template <typename T>
+class FixedAxisAngleErrorFunctionT
+    : public ConstraintErrorFunctionT<T, FixedAxisDataT<T>, 1, 1, 0> {
+ public:
+  explicit FixedAxisAngleErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, FixedAxisDataT<T>, 1, 1, 0>(skel, pt, lossAlpha, lossC) {}
+
+  explicit FixedAxisAngleErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : FixedAxisAngleErrorFunctionT<T>(
+            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-1f;
+
+ protected:
+  void evalFunction(
+      size_t constrIndex,
+      const JointStateT<T>& state,
+      Vector<T, 1>& f,
+      optional_ref<std::array<Vector3<T>, 1>> v = {},
+      optional_ref<std::array<Eigen::Matrix<T, 1, 3>, 1>> dfdv = {}) const final;
+};
+
+} // namespace momentum