pymomentum-cpu 0.0.0__cp312-cp312-manylinux_2_39_x86_64.whl

include/momentum/math/random.h

@@ -0,0 +1,310 @@
+/*
+ * 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/math/types.h>
+
+#include <random>
+
+namespace momentum {
+
+/// A consolidated random number generator to provide convenient APIs wrapping around the random
+/// library of C++ standard library <random>.
+///
+/// The default engine used in this class is std::mt19937, but it can be replaced, by specifying the
+/// template argument, with any preferable engines (see:
+/// https://en.cppreference.com/w/cpp/numeric/random).
+template <typename Generator_ = std::mt19937>
+class Random final {
+ public:
+  using Generator = Generator_;
+
+  /// Returns the singleton instance
+  [[nodiscard]] static Random& GetSingleton();
+
+  /// Constructor
+  explicit Random(uint32_t seed = std::random_device{}());
+
+  /// Generates a random scalar/vector/matrix from the uniform distribution
+  ///
+  /// The supported types are scalar, vector, or matrix where the elements can be either integer or
+  /// real. The supported integer types are whichever supported by std::uniform_int_distribution<T>
+  /// (see: https://en.cppreference.com/w/cpp/numeric/random/uniform_int_distribution), and the
+  /// supported real types are whichever std::is_floating_point<T>::value is true (see:
+  /// https://en.cppreference.com/w/cpp/types/is_floating_point).
+  ///
+  /// The generated random number is in the range of [min, max] for integer types or [min, max) for
+  /// real types. However, when built with the -ffast-math flag, the range for real types could be
+  /// [min, max] instead.
+  ///
+  /// @warning While the upper bound is typically exclusive for real types according to
+  /// std::uniform_real_distribution, there may be instances where the sampled value equals the
+  /// upper bound, depending on the platform and compiler specifics.
+  ///
+  /// For vector/matrix types, the bounds can be specified in either scalar or (the same dimension
+  /// of) vector/matrix. Use this function to specify the bounds in vector/matrix to apply different
+  /// bounds for each element in the generated vector/matrix. Use other versions of uniform() that
+  /// specify the bounds in scalar to apply the same bounds to all the elements of the generated
+  /// vector/matrix.
+  ///
+  /// @code
+  /// Random rand;
+  /// auto r0 = rand.uniform(0.0, 1.0); // random double in [0, 1)
+  /// auto r1 = rand.uniform(0.f, 1.f); // random float in [0, 1)
+  /// auto r3 = rand.uniform(0, 10);    // random int in [0, 10]
+  /// auto r4 = rand.uniform(Vector2f(0, 1), Vector2f(2, 3)); // random Vector2f in ([0, 2], [1, 3))
+  /// @endcode
+  ///
+  /// @tparam T The random number type to be generated
+  /// @param[in] min The lower bound of the random number
+  /// @param[in] max The upper bound of the random number
+  template <typename T>
+  [[nodiscard]] T uniform(const T& min, const T& max);
+
+  /// Generates a random vector/matrix from the uniform distribution, where the size is fixed
+  ///
+  /// This is a special case of T uniform(const T&, const T&) where the random number type is fixed
+  /// size of vector/matrix and the bounds are specified by the scalar type.
+  ///
+  /// @code
+  /// Random rand;
+  /// auto r0 = rand.uniform<Vector2f>(0, 1); // random Vector2f in ([0, 1], [0, 1))
+  /// @endcode
+  ///
+  /// @tparam FixedSizeT The random vector/matrix type to be generated
+  /// @param[in] min The lower bound in scalar.
+  /// @param[in] max The upper bound in scalar.
+  template <typename FixedSizeT>
+  [[nodiscard]] FixedSizeT uniform(
+      typename FixedSizeT::Scalar min,
+      typename FixedSizeT::Scalar max);
+
+  /// Generates a random vector from the uniform distribution, where the size is dynamic
+  ///
+  /// This is a special case of T uniform(const T&, const T&) where the random number type is
+  /// dynamic size of vector and the bounds are specified by the scalar type.
+  ///
+  /// @code
+  /// Random rand;
+  /// auto r0 = rand.uniform<VectorXf>(2, 0, 1); // random VectorXf of dimension 2 where each
+  /// element is in [0, 1)
+  /// @endcode
+  ///
+  /// @tparam DynamicVector The random vector type to be generated
+  /// @param[in] size The size of the vector.
+  /// @param[in] min The lower bound in scalar.
+  /// @param[in] max The upper bound in scalar.
+  template <typename DynamicVector>
+  [[nodiscard]] DynamicVector uniform(
+      Eigen::Index size,
+      typename DynamicVector::Scalar min,
+      typename DynamicVector::Scalar max);
+
+  /// Generates a random matrix from the uniform distribution, where the size is dynamic
+  ///
+  /// This is a special case of T uniform(const T&, const T&) where the random number type is
+  /// dynamic size of matrix and the bounds are specified by the scalar type.
+  ///
+  /// @code
+  /// Random rand;
+  /// auto r0 = rand.uniform<MatrixXf>(2, 3, 0, 1); // random MatrixXf of dimension 2x3 where each
+  /// element is in [0, 1)
+  /// @endcode
+  ///
+  /// @tparam DynamicMatrix The random matrix type to be generated
+  /// @param[in] rows The row size of the matrix.
+  /// @param[in] cols The column size of the matrix.
+  /// @param[in] min The lower bound in scalar.
+  /// @param[in] max The upper bound in scalar.
+  template <typename DynamicMatrix>
+  [[nodiscard]] DynamicMatrix uniform(
+      Eigen::Index rows,
+      Eigen::Index cols,
+      typename DynamicMatrix::Scalar min,
+      typename DynamicMatrix::Scalar max);
+
+  /// Generates a random quaternion from a uniform distribution on SO(3)
+  template <typename T>
+  [[nodiscard]] Quaternion<T> uniformQuaternion();
+
+  /// Generates a random rotation matrix from a uniform distribution on SO(3)
+  template <typename T>
+  [[nodiscard]] Matrix3<T> uniformRotationMatrix();
+
+  /// Generates a random isometry (rigid transformation) from a uniform distribution on SE(3)
+  /// based on input minimum and maximum 3D vectors for the translation part.
+  ///
+  /// @tparam Generator The type of random number generator.
+  /// @tparam T The type of isometry component, typically a float or double.
+  /// @param min The minimum 3D vector for the translation part.
+  /// @param max The maximum 3D vector for the translation part.
+  /// @return A random Isometry3 of type T.
+  template <typename T>
+  [[nodiscard]] Isometry3<T> uniformIsometry3(
+      const Vector3<T>& min = Vector3<T>::Zero(),
+      const Vector3<T>& max = Vector3<T>::Ones());
+
+  /// Generates a random affine transformation from a uniform distribution on the space of all
+  /// affine transformations.
+  ///
+  /// The transformation is created based on input minimum and maximum 3D vectors for the
+  /// translation part, and minimum and maximum scale factors for the linear part.
+  ///
+  /// @tparam Generator The type of random number generator.
+  /// @tparam T The type of affine transformation component, typically a float or double.
+  /// @param scaleMin The minimum scale factor for the linear part.
+  /// @param scaleMax The maximum scale factor for the linear part.
+  /// @param min The minimum 3D vector for the translation part.
+  /// @param max The maximum 3D vector for the translation part.
+  /// @return A random Affine3 of type T.
+  template <typename T>
+  [[nodiscard]] Affine3<T> uniformAffine3(
+      T scaleMin = 0.1,
+      T scaleMax = 10.0,
+      const Vector3<T>& min = Vector3<T>::Zero(),
+      const Vector3<T>& max = Vector3<T>::Ones());
+
+  /// Generates a random value from the Gaussian distribution
+  ///
+  /// @tparam T The random number type to be generated
+  /// @param[in] mean The mean of the Gaussian distribution
+  /// @param[in] sigma The standard deviation of the Gaussian distribution
+  template <typename T>
+  [[nodiscard]] T normal(const T& mean, const T& sigma);
+
+  /// Generates a random value from the Gaussian distribution
+  ///
+  /// @tparam FixedSizeT The random vector/matrix type to be generated
+  /// @param[in] mean The mean of the Gaussian distribution
+  /// @param[in] sigma The standard deviation of the Gaussian distribution
+  template <typename FixedSizeT>
+  [[nodiscard]] FixedSizeT normal(
+      typename FixedSizeT::Scalar mean,
+      typename FixedSizeT::Scalar sigma);
+
+  /// Generates a random value from the Gaussian distribution
+  ///
+  /// @tparam DynamicVector The random vector type to be generated
+  /// @param[in] size The size of the vector.
+  /// @param[in] mean The mean of the Gaussian distribution
+  /// @param[in] sigma The standard deviation of the Gaussian distribution
+  template <typename DynamicVector>
+  [[nodiscard]] DynamicVector normal(
+      Eigen::Index size,
+      typename DynamicVector::Scalar mean,
+      typename DynamicVector::Scalar sigma);
+
+  /// Generates a random value from the Gaussian distribution
+  ///
+  /// @tparam DynamicMatrix The random matrix type to be generated
+  /// @param[in] rows The row size of the matrix.
+  /// @param[in] cols The column size of the matrix.
+  /// @param[in] mean The mean of the Gaussian distribution
+  /// @param[in] sigma The standard deviation of the Gaussian distribution
+  template <typename DynamicMatrix>
+  [[nodiscard]] DynamicMatrix normal(
+      Eigen::Index rows,
+      Eigen::Index cols,
+      typename DynamicMatrix::Scalar mean,
+      typename DynamicMatrix::Scalar sigma);
+
+  /// Returns the seed.
+  [[nodiscard]] uint32_t getSeed() const;
+
+  /// Sets a new seed for the internal random number engine.
+  void setSeed(uint32_t seed);
+
+ private:
+  /// The seed used for the internal random number engine.
+  uint32_t seed_;
+
+  /// The internal random number engine.
+  Generator generator_;
+};
+
+/// Generates a random type T from the uniform distribution, using the global random number
+/// generator Random
+template <typename T>
+[[nodiscard]] T uniform(const T& min, const T& max);
+
+/// Generates a random fixed size vector/matrix from the uniform distribution, using the global
+/// random number generator Random
+template <typename FixedSizeT>
+[[nodiscard]] FixedSizeT uniform(typename FixedSizeT::Scalar min, typename FixedSizeT::Scalar max);
+
+/// Generates a random dynamic size vector from the uniform distribution, using the global random
+/// number generator Random
+template <typename DynamicVector>
+[[nodiscard]] DynamicVector
+uniform(Eigen::Index size, typename DynamicVector::Scalar min, typename DynamicVector::Scalar max);
+
+/// Generates a random type dynamic size matrix from the uniform distribution, using the global
+/// random number generator Random
+template <typename DynamicMatrix>
+[[nodiscard]] DynamicMatrix uniform(
+    Eigen::Index rows,
+    Eigen::Index cols,
+    typename DynamicMatrix::Scalar min,
+    typename DynamicMatrix::Scalar max);
+
+/// Generates a random quaternion from a uniform distribution on SO(3)
+template <typename T>
+[[nodiscard]] Quaternion<T> uniformQuaternion();
+
+/// Generates a random rotation matrix from a uniform distribution on SO(3)
+template <typename T>
+[[nodiscard]] Matrix3<T> uniformRotationMatrix();
+
+/// Generates a random isometry (rigid transformation) from a uniform distribution on SE(3)
+/// based on input minimum and maximum 3D vectors for the translation part.
+template <typename T>
+[[nodiscard]] Isometry3<T> uniformIsometry3(
+    const Vector3<T>& min = Vector3<T>::Zero(),
+    const Vector3<T>& max = Vector3<T>::Ones());
+
+/// Generates a random affine transformation from a uniform distribution on the space of all
+/// affine transformations.
+template <typename T>
+[[nodiscard]] Affine3<T> uniformAffine3(
+    T scaleMin = 0.1,
+    T scaleMax = 10.0,
+    const Vector3<T>& min = Vector3<T>::Zero(),
+    const Vector3<T>& max = Vector3<T>::Ones());
+
+/// Generates a random type T from the Gaussian distribution, using the global random number
+/// generator Random
+template <typename T>
+[[nodiscard]] T normal(const T& mean, const T& sigma);
+
+/// Generates a random fixed size vector/matrix from the Gaussian distribution, using the global
+/// random number generator Random
+template <typename FixedSizeT>
+[[nodiscard]] FixedSizeT normal(
+    typename FixedSizeT::Scalar mean,
+    typename FixedSizeT::Scalar sigma);
+
+/// Generates a random dynamic size vector from the Gaussian distribution, using the global random
+/// number generator Random
+template <typename DynamicVector>
+[[nodiscard]] DynamicVector normal(
+    Eigen::Index size,
+    typename DynamicVector::Scalar mean,
+    typename DynamicVector::Scalar sigma);
+
+/// Generates a random dynamic size matrix from the Gaussian distribution, using the global random
+/// number generator Random
+template <typename DynamicMatrix>
+[[nodiscard]] DynamicMatrix normal(
+    Eigen::Index rows,
+    Eigen::Index cols,
+    typename DynamicMatrix::Scalar mean,
+    typename DynamicMatrix::Scalar sigma);
+
+} // namespace momentum
+
+#include <momentum/math/random-inl.h>

include/momentum/math/simd_generalized_loss.h

@@ -0,0 +1,40 @@
+/*
+ * 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/math/generalized_loss.h>
+#include <momentum/simd/simd.h>
+
+namespace momentum {
+
+/// SIMD version of the generalized loss function.
+template <typename T>
+class SimdGeneralizedLossT : public GeneralizedLossT<T> {
+ public:
+  using Base = GeneralizedLossT<T>;
+
+  /// Constructs a SIMD generalized loss function
+  ///
+  /// @param a Alpha parameter controlling the shape of the loss function
+  /// @param c Scale parameter controlling the width of the quadratic region
+  explicit SimdGeneralizedLossT(const T& a = Base::kL2, const T& c = T(1));
+
+  /// Computes the loss value for a given squared error
+  ///
+  /// @param sqrError Squared error term (SIMD packet)
+  /// @return Loss value as a SIMD packet
+  [[nodiscard]] Packet<T> value(const Packet<T>& sqrError) const;
+
+  /// Computes the derivative of the loss function
+  ///
+  /// @param sqrError Squared error term (SIMD packet)
+  /// @return Derivative value as a SIMD packet
+  [[nodiscard]] Packet<T> deriv(const Packet<T>& sqrError) const;
+};
+
+} // namespace momentum

include/momentum/math/transform.h

@@ -0,0 +1,229 @@
+/*
+ * 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/math/types.h>
+
+#include <Eigen/Core>
+#include <Eigen/Geometry>
+
+namespace momentum {
+
+/// Lightweight transform with separate rotation, translation, and uniform scaling
+///
+/// This class provides a more efficient alternative to Eigen's Affine3 for transformations
+/// that only involve rotation, translation, and uniform scaling. It stores these components
+/// separately, which allows for more efficient operations and memory usage.
+///
+/// @tparam T The scalar type (float or double)
+template <typename T>
+struct TransformT {
+  Quaternion<T> rotation;
+  Vector3<T> translation;
+  T scale;
+
+  /// Creates a transform with only rotation
+  [[nodiscard]] static TransformT<T> makeRotation(const Quaternion<T>& rotation_in);
+
+  /// Creates a transform with only translation
+  [[nodiscard]] static TransformT<T> makeTranslation(const Vector3<T>& translation_in);
+
+  /// Creates a transform with only scaling
+  [[nodiscard]] static TransformT<T> makeScale(const T& scale_in);
+
+  /// Creates a random transform
+  ///
+  /// @param[in] translation If true, use random translation (-1 to 1)
+  /// @param[in] rotation If true, use random rotation
+  /// @param[in] scale If true, use random scale (0.5 to 2.0)
+  [[nodiscard]] static TransformT<T>
+  makeRandom(bool translation = true, bool rotation = true, bool scale = true);
+
+  /// Converts from Eigen's Affine3
+  [[nodiscard]] static TransformT<T> fromAffine3(const Affine3<T>& other);
+
+  /// Converts from 4x4 matrix
+  [[nodiscard]] static TransformT<T> fromMatrix(const Matrix4<T>& other);
+
+  /// Creates identity transform
+  TransformT() : rotation(Quaternion<T>::Identity()), translation(Vector3<T>::Zero()), scale(1) {
+    // Empty
+  }
+
+  /// Copy constructor with type conversion
+  ///
+  /// @tparam T2 Source scalar type
+  template <typename T2>
+  explicit TransformT(const TransformT<T2>& other)
+      : rotation(other.rotation.template cast<T>()),
+        translation(other.translation.template cast<T>()),
+        scale(other.scale) {
+    // Empty
+  }
+
+  /// Constructor with components
+  explicit TransformT(
+      const Vector3<T>& translation_in,
+      const Quaternion<T>& rotation_in = Quaternion<T>::Identity(),
+      const T& scale_in = T(1))
+      : rotation(rotation_in), translation(translation_in), scale(scale_in) {
+    // Empty
+  }
+
+  /// Move constructor
+  explicit TransformT(
+      Vector3<T>&& translation_in,
+      Quaternion<T>&& rotation_in = Quaternion<T>::Identity(),
+      T&& scale_in = T(1))
+      : rotation(std::move(rotation_in)),
+        translation(std::move(translation_in)),
+        scale(std::move(scale_in)) {
+    // Empty
+  }
+
+  /// Constructor from Affine3
+  explicit TransformT(const Affine3<T>& other) {
+    *this = fromAffine3(other);
+  }
+
+  /// Constructor from 4x4 matrix
+  explicit TransformT(const Matrix4<T>& other) {
+    *this = fromMatrix(other);
+  }
+
+  /// Assignment from Affine3
+  TransformT<T>& operator=(const Affine3<T>& other) {
+    *this = fromAffine3(other);
+    return *this;
+  }
+
+  /// Assignment from 4x4 matrix
+  TransformT<T>& operator=(const Matrix4<T>& other) {
+    *this = fromMatrix(other);
+    return *this;
+  }
+
+  /// Combines transforms (applies other first, then this)
+  [[nodiscard]] TransformT<T> operator*(const TransformT<T>& other) const {
+    // [ s_1*R_1 t_1 ] * [ s_2*R_2 t_2 ] = [ s_1*s_2*R_1*R_2  s_1*R_1*t_2 + t_1 ]
+    // [     0    1  ]   [     0    1  ]   [        0                     1     ]
+    const Vector3<T> trans = translation + rotation * (scale * other.translation);
+    const Quaternion<T> rot = rotation * other.rotation;
+    const T newScale = scale * other.scale;
+    return TransformT<T>(std::move(trans), std::move(rot), std::move(newScale));
+  }
+
+  /// Multiplies with Affine3
+  [[nodiscard]] Affine3<T> operator*(const Affine3<T>& other) const {
+    Affine3<T> out = Affine3<T>::Identity();
+    out.linear().noalias() = toLinear() * other.linear();
+    out.translation().noalias() = rotation * (scale * other.translation()) + translation;
+    return out;
+  }
+
+  /// Transforms a point
+  [[nodiscard]] Vector3<T> operator*(const Vector3<T>& other) const {
+    return transformPoint(other);
+  }
+
+  /// Converts to Affine3
+  explicit operator Affine3<T>() const {
+    return toAffine3();
+  }
+
+  /// Converts to Affine3
+  [[nodiscard]] Affine3<T> toAffine3() const;
+
+  /// Converts to 4x4 matrix
+  [[nodiscard]] Matrix4<T> toMatrix() const {
+    Matrix4<T> out = Matrix4<T>::Zero();
+    out.template topLeftCorner<3, 3>().noalias() = rotation.toRotationMatrix() * scale;
+    out.template topRightCorner<3, 1>() = translation;
+    out(3, 3) = T(1);
+    return out;
+  }
+
+  /// Returns rotation as 3x3 matrix
+  [[nodiscard]] Matrix3<T> toRotationMatrix() const {
+    return rotation.toRotationMatrix();
+  }
+
+  /// Returns rotation*scale as 3x3 matrix
+  [[nodiscard]] Matrix3<T> toLinear() const {
+    return toRotationMatrix() * scale;
+  }
+
+  /// Gets the X axis of the rotation
+  [[nodiscard]] Vector3<T> getAxisX() const {
+    // Optimized version of Eigen::Quaternion::toRotationMatrix()
+    Vector3<T> axis;
+
+    const T ty = T(2) * rotation.y();
+    const T tz = T(2) * rotation.z();
+    const T twy = ty * rotation.w();
+    const T twz = tz * rotation.w();
+    const T txy = ty * rotation.x();
+    const T txz = tz * rotation.x();
+    const T tyy = ty * rotation.y();
+    const T tzz = tz * rotation.z();
+
+    axis[0] = T(1) - (tyy + tzz);
+    axis[1] = txy + twz;
+    axis[2] = txz - twy;
+
+    return axis;
+  }
+
+  /// Applies full transform to a point
+  [[nodiscard]] Vector3<T> transformPoint(const Vector3<T>& pt) const {
+    return translation + rotation * (scale * pt).eval();
+  }
+
+  /// Applies only rotation to a vector
+  [[nodiscard]] Vector3<T> rotate(const Vector3<T>& vec) const;
+
+  /// Computes inverse transform
+  [[nodiscard]] TransformT<T> inverse() const;
+
+  /// Converts to different scalar type
+  ///
+  /// @tparam T2 Target scalar type
+  template <typename T2>
+  [[nodiscard]] TransformT<T2> cast() const {
+    return TransformT<T2>(
+        this->translation.template cast<T2>(),
+        this->rotation.template cast<T2>(),
+        static_cast<T2>(this->scale));
+  }
+
+  /// Checks if this transform is approximately equal to another
+  [[nodiscard]] bool isApprox(
+      const TransformT<T>& other,
+      T tol = Eigen::NumTraits<T>::dummy_precision()) const {
+    // TODO: Improve
+    return toAffine3().isApprox(other.toAffine3(), tol);
+  }
+};
+
+template <typename T>
+using TransformListT =
+    std::vector<TransformT<T>>; // structure describing a the state of all joints in a skeleton
+
+template <typename T>
+TransformT<T> blendTransforms(
+    gsl::span<const TransformT<T>> transforms,
+    gsl::span<const T> weights);
+
+/// Spherical linear interpolation between two transforms
+template <typename T>
+TransformT<T> slerp(const TransformT<T>& t1, const TransformT<T>& t2, T weight);
+
+using Transform = TransformT<float>;
+using TransformList = TransformListT<float>;
+
+} // namespace momentum