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

include/momentum/character_sequence_solver/sequence_solver_function.h

@@ -0,0 +1,134 @@
+/*
+ * 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 {
+
+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_;
+  }
+
+  [[nodiscard]] bool needsMesh() const {
+    return needsMesh_;
+  }
+
+  // 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 states_.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(
+      gsl::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);
+
+ private:
+  const Character& character_;
+  const ParameterTransformT<T>& parameterTransform_;
+  std::vector<SkeletonStateT<T>> states_;
+  std::vector<MeshStateT<T>> meshStates_;
+  VectorX<bool> activeJointParams_;
+
+  std::vector<ModelParametersT<T>> frameParameters_;
+
+  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;
+
+  bool needsMesh_ = false;
+
+  friend class SequenceSolverT<T>;
+};
+
+} // namespace momentum

include/momentum/character_sequence_solver/state_sequence_error_function.h

@@ -0,0 +1,109 @@
+/*
+ * 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>
+
+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);
+  explicit StateSequenceErrorFunctionT(const Character& character);
+
+  [[nodiscard]] size_t numFrames() const final {
+    return 2;
+  }
+
+  double getError(
+      gsl::span<const ModelParametersT<T>> modelParameters,
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      gsl::span<const MeshStateT<T>> /* meshStates */) const final;
+  double getGradient(
+      gsl::span<const ModelParametersT<T>> modelParameters,
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      gsl::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(
+      gsl::span<const ModelParametersT<T>> modelParameters,
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      gsl::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_;
+
+ public:
+  // weights for the error functions
+  static constexpr T kPositionWeight = 1e-3f;
+  static constexpr T kOrientationWeight = 1e+0f;
+};
+
+} // namespace momentum

include/momentum/character_sequence_solver/vertex_sequence_error_function.h

@@ -0,0 +1,128 @@
+/*
+ * 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/skeleton_state.h>
+#include <momentum/character_sequence_solver/fwd.h>
+#include <momentum/character_sequence_solver/sequence_error_function.h>
+#include <momentum/character_solver/vertex_error_function.h>
+#include <momentum/math/mesh.h>
+
+#include <gsl/span>
+
+namespace momentum {
+
+/// Constraint structure for vertex velocity.
+template <typename T>
+struct VertexVelocityConstraintT {
+  int vertexIndex = -1;
+  T weight = 1;
+  Eigen::Vector3<T> targetVelocity;
+
+  template <typename T2>
+  VertexVelocityConstraintT<T2> cast() const {
+    return {
+        this->vertexIndex, static_cast<T2>(this->weight), this->targetVelocity.template cast<T2>()};
+  }
+};
+
+/// Error function that penalizes differences in vertex velocity between source and target motion.
+/// This function computes vertex velocities by taking the difference between consecutive frames
+/// and penalizes deviations from target vertex velocities. It combines the sequence aspect of
+/// StateSequenceErrorFunction with vertex constraints from VertexErrorFunction.
+template <typename T>
+class VertexSequenceErrorFunctionT : public SequenceErrorFunctionT<T> {
+ public:
+  explicit VertexSequenceErrorFunctionT(const Character& character);
+
+  [[nodiscard]] size_t numFrames() const final {
+    return 2;
+  }
+
+  double getError(
+      gsl::span<const ModelParametersT<T>> modelParameters,
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      gsl::span<const MeshStateT<T>> meshStates) const final;
+  double getGradient(
+      gsl::span<const ModelParametersT<T>> modelParameters,
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      gsl::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(
+      gsl::span<const ModelParametersT<T>> modelParameters,
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      gsl::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;
+
+  /// Add a vertex velocity constraint.
+  /// @param vertexIndex Index of the vertex to constrain.
+  /// @param weight Weight for this constraint.
+  /// @param targetVelocity Target velocity for the vertex.
+  void addConstraint(int vertexIndex, T weight, const Eigen::Vector3<T>& targetVelocity);
+
+  /// Clear all vertex velocity constraints.
+  void clearConstraints();
+
+  /// Get all vertex velocity constraints.
+  [[nodiscard]] const std::vector<VertexVelocityConstraintT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+  /// Get the number of constraints.
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+
+  /// Get the character reference.
+  [[nodiscard]] const Character& getCharacter() const {
+    return character_;
+  }
+
+  static constexpr T kVelocityWeight = 1e-3f;
+
+ private:
+  /// Calculate gradient for a single vertex velocity constraint.
+  double calculateVelocityGradient(
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      const VertexVelocityConstraintT<T>& constraint,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  /// Calculate Jacobian for a single vertex velocity constraint.
+  double calculateVelocityJacobian(
+      gsl::span<const SkeletonStateT<T>> skelStates,
+      const VertexVelocityConstraintT<T>& constraint,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      Eigen::Index startRow) const;
+
+  /// Update meshes for both frames.
+  void updateMeshes(
+      gsl::span<const ModelParametersT<T>> modelParameters,
+      gsl::span<const SkeletonStateT<T>> skelStates) const;
+
+  const Character& character_;
+  std::vector<VertexVelocityConstraintT<T>> constraints_;
+
+  // Meshes for the two frames
+  mutable std::unique_ptr<MeshT<T>> neutralMesh_;
+  mutable std::unique_ptr<MeshT<T>> restMesh_; // Single rest mesh used for both frames
+  mutable std::unique_ptr<MeshT<T>> posedMesh0_;
+  mutable std::unique_ptr<MeshT<T>> posedMesh1_;
+};
+
+} // namespace momentum

include/momentum/character_solver/aim_error_function.h

@@ -0,0 +1,112 @@
+/*
+ * 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 {
+
+/// Aim constraint data: a local ray (point, dir) passes a global point target
+template <typename T>
+struct AimDataT : ConstraintData {
+  /// The origin of the local ray
+  Vector3<T> localPoint;
+  /// The direction of the local ray
+  Vector3<T> localDir;
+  /// The global aim target
+  Vector3<T> globalTarget;
+
+  explicit AimDataT(
+      const Vector3<T>& inLocalPt,
+      const Vector3<T>& inLocalDir,
+      const Vector3<T>& inTarget,
+      size_t pIndex,
+      float w,
+      const std::string& n = {})
+      : ConstraintData(pIndex, w, n),
+        localPoint(inLocalPt),
+        localDir(inLocalDir.normalized()),
+        globalTarget(inTarget) {}
+};
+
+/// The AimDistErrorFunction computes the distance from the aim target point to an aim ray (an
+/// original and a positive direction). The distance is minimized if the target lies on the ray (ie.
+/// on the positive direction), or when the target point is close to the ray origin. Another way to
+/// look at this loss is to minimize the element-wise differences between the aim vector's
+/// projection to the ray direction and the aim vector itself (therefore a shorter aim vector also
+/// has a smaller norm).
+template <typename T>
+class AimDistErrorFunctionT : public ConstraintErrorFunctionT<T, AimDataT<T>, 3, 2, 1> {
+ public:
+  explicit AimDistErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, AimDataT<T>, 3, 2, 1>(skel, pt, lossAlpha, lossC) {}
+
+  explicit AimDistErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : AimDistErrorFunctionT<T>(
+            character.skeleton,
+            character.parameterTransform,
+            lossAlpha,
+            lossC) {}
+
+  /// Default constant weight in MarkerErrorFunction (same as fixed axis). 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>, 2>> v = {},
+      optional_ref<std::array<Eigen::Matrix3<T>, 2>> dfdv = {}) const final;
+};
+
+/// The AimDirErrorFunction computes the element-wise differences between the normalized aim vector
+/// and the ray direction. Because the aim vector is normalized, the distance to the ray origin
+/// doesn't affect the loss as in AimDistErrorFunction. When the aim vector is close to zero length,
+/// we will clamp it to zero.
+template <typename T>
+class AimDirErrorFunctionT : public ConstraintErrorFunctionT<T, AimDataT<T>, 3, 2, 1> {
+ public:
+  explicit AimDirErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, AimDataT<T>, 3, 2, 1>(skel, pt, lossAlpha, lossC) {}
+
+  explicit AimDirErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : AimDirErrorFunctionT<T>(
+            character.skeleton,
+            character.parameterTransform,
+            lossAlpha,
+            lossC) {}
+
+  /// Default constant weight in MarkerErrorFunction (same as fixed axis). 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>, 2>> v = {},
+      optional_ref<std::array<Eigen::Matrix3<T>, 2>> dfdv = {}) const final;
+};
+} // namespace momentum

include/momentum/character_solver/collision_error_function.h

@@ -0,0 +1,92 @@
+/*
+ * 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 <axel/Bvh.h>
+
+#include <unordered_set>
+#include <vector>
+
+namespace momentum {
+
+template <typename T>
+class CollisionErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  explicit CollisionErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const CollisionGeometry& cg);
+
+  explicit CollisionErrorFunctionT(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;
+
+  [[nodiscard]] std::vector<Vector2i> getCollisionPairs() const;
+
+ protected:
+  void updateCollisionPairs();
+
+  // Update collisionState_ and bvh_ 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_;
+
+  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/collision_error_function_stateless.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/collision_geometry.h>
+#include <momentum/character/collision_geometry_state.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+
+namespace momentum {
+
+/// NOTE: This is a resurrected version of the `CollisionErrorFunction` (previously updated by
+/// D50818775). Its stateless nature makes it ideal for testing in multi-threaded settings.
+///
+/// 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 CollisionErrorFunctionStatelessT : public SkeletonErrorFunctionT<T> {
+ public:
+  explicit CollisionErrorFunctionStatelessT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const CollisionGeometry& cg);
+
+  explicit CollisionErrorFunctionStatelessT(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;
+
+  [[nodiscard]] std::vector<Vector2i> getCollisionPairs() const;
+
+ protected:
+  void updateCollisionPairs();
+
+  const CollisionGeometry collisionGeometry;
+
+  // [0] and [1] index the collision geometry, [2] indexes the joint after
+  // which we can stop comparing as it's parent to both
+  std::vector<Vector3i> collisionPairs;
+
+  VectorX<bool> collisionActive;
+
+  CollisionGeometryStateT<T> collisionState;
+
+  // weights for the error functions
+  static constexpr T kCollisionWeight = 5e-3f;
+};
+
+} // namespace momentum