pymomentum-cpu 0.1.78.post11__cp312-cp312-win_amd64.whl

include/axel/BvhBase.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 <array>
+#include <vector>
+
+#include "axel/BoundingBox.h"
+#include "axel/common/Types.h"
+
+namespace axel {
+
+template <typename S>
+class BvhBase {
+ public:
+  using Scalar = S;
+  using QueryBuffer = std::array<uint32_t, 512>;
+  BvhBase() = default;
+  virtual ~BvhBase() = default;
+
+  BvhBase(const BvhBase&) = default;
+  BvhBase& operator=(const BvhBase&) = default;
+
+  BvhBase(BvhBase&&) noexcept = default;
+  BvhBase& operator=(BvhBase&&) noexcept = default;
+
+  /**
+   * @brief Rebuilds the bvh with a new vector of bounding boxes.
+   * @param bboxes The new bounding boxes.
+   */
+  virtual void setBoundingBoxes(const std::vector<BoundingBox<S>>& bboxes) = 0;
+
+  /**
+   * @brief Checks for an intersection between a given bounding box and the bvh.
+   * All IDs of the primitives that overlap with the given bounding box will be returned
+   * in a dynamically-allocated vector.
+   * @param box The input bounding box to intersect with the bvh.
+   * @return A vector of primitive IDs that intersect with the box.
+   */
+  [[nodiscard]] virtual std::vector<uint32_t> query(const BoundingBox<S>& box) const = 0;
+
+  /**
+   * @brief Checks for an intersection between a given bounding box and the bvh.
+   * The query permits only up to first 512 intersections, which is the limit
+   * of the query buffer. In practice this shouldn't make a difference,
+   * but performance-wise we don't allocate tens of thousands of vectors on the heap.
+   * @param box The input bounding box to perform the query with.
+   * @param hits The buffer where the hit primitive IDs will be stored.
+   * @return The number of hit primitive IDs. When using the hits buffer afterwards,
+   * hit primitives are within [0, hitCount - 1] interval of the hits buffer.
+   */
+  virtual uint32_t query(const BoundingBox<S>& box, QueryBuffer& hits) const = 0;
+
+  /**
+   * @brief Checks for an intersection between a ray and the bvh.
+   * @param origin The origin of the ray in the same space as the bvh data.
+   * @param direction The direction of the ray in the same space as the bvh data.
+   * @return A vector of primitive IDs that intersect with the ray.
+   */
+  [[nodiscard]] virtual std::vector<uint32_t> query(
+      const Eigen::Vector3<S>& origin,
+      const Eigen::Vector3<S>& direction) const = 0;
+
+  [[nodiscard]] virtual Size getPrimitiveCount() const = 0;
+};
+
+using BvhBasef = BvhBase<float>;
+using BvhBased = BvhBase<double>;
+
+} // namespace axel

include/axel/BvhCommon.h

@@ -0,0 +1,43 @@
+/*
+ * 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 <optional>
+
+#include <Eigen/Core>
+
+namespace axel {
+
+template <typename S>
+struct IntersectionResult {
+  using Scalar = S;
+  int32_t triangleId{0};
+  S hitDistance{std::numeric_limits<S>::max()};
+  Eigen::Vector3<S> hitPoint;
+  Eigen::Vector3<S> baryCoords;
+};
+
+using IntersectionResultf = IntersectionResult<float>;
+using IntersectionResultd = IntersectionResult<double>;
+
+// TODO(nemanjab): Fix-it: unify index types.
+inline constexpr uint32_t kInvalidTriangleIdx = std::numeric_limits<uint32_t>::max();
+
+template <typename S>
+struct ClosestSurfacePointResult {
+  Eigen::Vector3<S> point;
+  uint32_t triangleIdx = kInvalidTriangleIdx;
+
+  // TODO: Remove optional once all the queries support barycentric coordinates.
+  std::optional<Eigen::Vector3<S>> baryCoords;
+};
+
+using ClosestSurfacePointResultf = ClosestSurfacePointResult<float>;
+using ClosestSurfacePointResultd = ClosestSurfacePointResult<double>;
+
+} // namespace axel

include/axel/BvhEmbree.h

@@ -0,0 +1,86 @@
+/*
+ * 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 <array>
+#include <optional>
+#include <vector>
+
+#include <embree4/rtcore.h>
+
+#include "axel/BoundingBox.h"
+#include "axel/BvhBase.h"
+
+namespace axel {
+
+/**
+ * @brief Encapsulates an Embree BVH for a single geometry.
+ *
+ * The data structure can be initialized/updated with setBoundingBoxes(...)
+ * or preferably with a vector of RTCBoundPrimitive structs through setPrimitives(...).
+ *
+ * The BVH supports bounding boxes and ray queries to conform to the interface
+ * of the other BVH implementations in this library.
+ */
+class BVHEmbree final : public BvhBased {
+ public:
+  struct alignas(16) Node {
+    // A bounding box with SIMD-friendly storage.
+    std::array<RTCBounds, 2> bounds{};
+    std::array<BVHEmbree::Node*, 2> children{nullptr, nullptr};
+    uint32_t primID{0};
+    bool isLeaf{false};
+
+    Node() = default;
+    explicit Node(const RTCBuildPrimitive& bounds);
+  };
+
+  /**
+   * @brief Construct a new, empty BVH based on Embree.
+   */
+  explicit BVHEmbree(std::optional<uint32_t> embreeThreadCount = std::nullopt);
+  ~BVHEmbree() override;
+
+  BVHEmbree(const BVHEmbree& other) = delete;
+  BVHEmbree(BVHEmbree&& other) noexcept;
+  BVHEmbree& operator=(const BVHEmbree& other) = delete;
+  BVHEmbree& operator=(BVHEmbree&& other) noexcept;
+
+  void setBoundingBoxes(const std::vector<BoundingBoxd>& bboxes) override;
+
+  /**
+   * @brief Updates the BVH with a new vector of bounding boxes.
+   * The input parameter can't be made const because Embree takes a non-const ptr
+   * to its data.
+   * @param prims The new bounding boxes, in Embree's expected, aligned structs.
+   */
+  void setPrimitives(std::vector<RTCBuildPrimitive>& prims);
+
+  uint32_t query(const BoundingBoxd& box, QueryBuffer& hits) const override;
+  std::vector<unsigned int> query(const BoundingBoxd& box) const override;
+  std::vector<unsigned int> query(const Eigen::Vector3d& origin, const Eigen::Vector3d& direction)
+      const override;
+
+  Size getPrimitiveCount() const override;
+
+ private:
+  // Manually ref-counted handle to the device used as Embree's context.
+  RTCDevice device_{nullptr};
+
+  // Struct for BVH build parameters.
+  RTCBuildArguments buildArguments_;
+
+  // Owning handle to the BVH handle.
+  RTCBVH bvh_{nullptr};
+
+  // This is a non-owning pointer to Embree's internal thread-local memory allocation.
+  // All the memory will be freed once we call rtcReleaseBVH(bvh_) in the destructor of this class.
+  Node* root_{nullptr};
+};
+
+} // namespace axel

include/axel/BvhFactory.h

@@ -0,0 +1,34 @@
+/*
+ * 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 <memory>
+#include <optional>
+#include <string_view>
+
+#include "axel/BvhBase.h"
+
+namespace axel {
+
+/**
+ * @brief Creates a bvh that speeds up spatial look-up for collision
+ * detection.
+ * @param type Type of the acceleration structure. Supported values are:
+ * - "default": instantiates an implementation-specific default BVH
+ * - "embree": instantiates an Embree-based BVH on supported (desktop) platforms. Otherwise,
+ * creates a "default" BVH.
+ * - "fast-bvh": instantiates a BVH based on the https://github.com/brandonpelfrey/Fast-BVH.
+ * @param threadCount Specify how many threads to use for a given acceleration structure if
+ * supported.
+ * @return A unique pointer to the created acceleration structure.
+ */
+std::unique_ptr<BvhBased> createBvh(
+    std::string_view type = "default",
+    std::optional<uint32_t> threadCount = std::nullopt);
+
+} // namespace axel

include/axel/Checks.h

@@ -0,0 +1,21 @@
+/*
+ * 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
+
+#if defined(AXEL_WITH_XR_LOGGER)
+
+#include <logging/Checks.h>
+
+#else
+
+#include <cassert>
+
+// TODO: Support asserts with messages as XR_CHECK does
+#define XR_CHECK(condition, ...) assert(condition)
+
+#endif

include/axel/DualContouring.h

@@ -0,0 +1,79 @@
+/*
+ * 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 <array>
+#include <unordered_map>
+#include <vector>
+
+#include <Eigen/Core>
+#include <Eigen/Dense>
+#include <gsl/span>
+
+#include "axel/BoundingBox.h"
+#include "axel/SignedDistanceField.h"
+#include "axel/common/Types.h"
+
+namespace axel {
+
+/**
+ * Result of dual contouring operation.
+ * Dual contouring naturally produces quads, so this always contains quads.
+ * Use triangulateQuads() to convert to triangles if needed.
+ */
+template <typename S>
+struct DualContouringResult {
+  using Scalar = S;
+
+  /// Generated vertices
+  std::vector<Eigen::Vector3<S>> vertices;
+
+  /// Generated quads (dual contouring naturally produces quads)
+  std::vector<Eigen::Vector4i> quads;
+
+  /// Success flag
+  bool success = false;
+
+  /// Number of cells processed
+  size_t processedCells = 0;
+
+  /// Number of vertices generated
+  size_t generatedVertices = 0;
+};
+
+/**
+ * Extract an isosurface from a signed distance field using dual contouring.
+ *
+ * Dual contouring places vertices inside grid cells (rather than on edges like marching cubes)
+ * and uses both function values and gradients to determine optimal vertex placement.
+ * This results in better preservation of sharp features and corners compared to marching cubes.
+ *
+ * The algorithm works by:
+ * 1. Finding all cells that intersect the isosurface (sign changes across cell corners)
+ * 2. Placing one vertex at each intersecting cell, positioned on the surface using gradient descent
+ * 3. Generating quads for each edge crossing that connects 4 adjacent cells
+ *
+ * @param sdf The signed distance field to extract from
+ * @param isovalue The isovalue to extract (typically 0.0 for zero level set)
+ * @return Result containing extracted mesh
+ */
+template <typename ScalarType>
+DualContouringResult<ScalarType> dualContouring(
+    const SignedDistanceField<ScalarType>& sdf,
+    ScalarType isovalue = ScalarType{0.0});
+
+/**
+ * Triangulate a quad mesh into triangles.
+ * Each quad is split into two triangles using the diagonal (0,2).
+ *
+ * @param quads Vector of quads to triangulate
+ * @return Vector of triangles
+ */
+std::vector<Eigen::Vector3i> triangulateQuads(const std::vector<Eigen::Vector4i>& quads);
+
+} // namespace axel

include/axel/KdTree.h

@@ -0,0 +1,199 @@
+/*
+ * 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 <optional>
+
+#include <nanoflann/nanoflann.hpp>
+#include <Eigen/Core>
+#include <Eigen/Eigen>
+
+namespace axel {
+
+namespace detail {
+
+template <typename EigenType>
+inline constexpr bool IsEigenTypeWithStorage =
+    std::is_base_of_v<Eigen::PlainObjectBase<std::decay_t<EigenType>>, std::decay_t<EigenType>>;
+
+template <typename EigenType, Eigen::Index Length>
+inline constexpr bool IsContiguousEigenVectorBaseWithLength =
+    (std::decay_t<EigenType>::RowsAtCompileTime == Length ||
+     std::decay_t<EigenType>::ColsAtCompileTime == Length) &&
+    (std::decay_t<EigenType>::RowsAtCompileTime == 1 ||
+     std::decay_t<EigenType>::ColsAtCompileTime == 1) &&
+    std::decay_t<EigenType>::InnerStrideAtCompileTime == 1; // Make sure the memory is contiguous.
+} // namespace detail
+
+enum class SortByDistance { True, False };
+
+template <typename EigenMatrixType>
+class KdTree {
+ public:
+  // The scalar type of the point coordinates.
+  using Scalar = typename EigenMatrixType::Scalar;
+
+  // The storage of data needs to be done with an Eigen::Matrix.
+  static_assert(
+      detail::IsEigenTypeWithStorage<EigenMatrixType>,
+      "KdTree's current implementation only works with Eigen matrices.");
+
+  // Dimensionality of the tree is determined by compile-time columns in the point matrix.
+  // It can also be dynamic, i.e. Eigen::Dynamic.
+  static constexpr Eigen::Index kDimensionality = EigenMatrixType::ColsAtCompileTime;
+
+  using TreeAdaptorType = nanoflann::KDTreeEigenMatrixAdaptor<EigenMatrixType, kDimensionality>;
+
+  // The underlying index for elements, i.e. Eigen::Index.
+  using IndexType = typename TreeAdaptorType::IndexType;
+
+  using RadiusSearchResults = std::vector<std::pair<IndexType, Scalar>>;
+
+  struct KnnSearchResults {
+    std::vector<IndexType> indices;
+    std::vector<Scalar> squaredDistances;
+  };
+
+  struct ClosestSearchResult {
+    IndexType index{};
+    Scalar squaredDistance{};
+  };
+
+  // Only allow lvalue and rvalue references of points that are exactly the Eigen::Matrix type
+  // we used to construct the tree.
+  template <
+      typename T,
+      typename = std::enable_if_t<
+          detail::IsEigenTypeWithStorage<T> && // This check might be obsolete due to transitivity.
+          std::is_same_v<std::decay_t<T>, std::decay_t<EigenMatrixType>>>>
+  explicit KdTree(T&& points) : points_(std::forward<T>(points)), tree_(points_.cols(), points_) {
+    tree_.index->buildIndex();
+  }
+
+  // Returns the number of points in the tree.
+  size_t getSize() const {
+    return static_cast<size_t>(tree_.index->size());
+  }
+
+  bool isEmpty() const {
+    return tree_.index->size() == 0;
+  }
+
+  /**
+   * @brief Returns all nearby points in the tree to a given query point with the given radius.
+   *
+   * @param query The point to query the tree with.
+   * @param radius The radius to search with.
+   * @param sort If SortByDistance::True, the results are sorted by distance.
+   * @param results The results contain the indices and squared distances of the points within the
+   * given radius.
+   */
+  template <typename EigenVectorType>
+  void radiusSearch(
+      EigenVectorType&& query,
+      const Scalar radius,
+      const SortByDistance sort,
+      RadiusSearchResults& results) const {
+    static_assert(detail::IsContiguousEigenVectorBaseWithLength<EigenVectorType, kDimensionality>);
+    nanoflann::SearchParams params{};
+    params.sorted = sort == SortByDistance::True;
+    tree_.index->radiusSearch(query.data(), radius * radius, results, params);
+  }
+
+  /**
+   * @brief Returns all nearby points in the tree to a given query point with the given radius.
+   *
+   * @param query The point to query the tree with.
+   * @param radius The radius to search with.
+   * @param sort If SortByDistance::True, the results are sorted by distance.
+   * @return The results contain the indices and squared distances of the points within the given
+   * radius.
+   */
+  template <typename EigenVectorType>
+  RadiusSearchResults
+  radiusSearch(EigenVectorType&& query, const Scalar radius, const SortByDistance sort) const {
+    static_assert(detail::IsContiguousEigenVectorBaseWithLength<EigenVectorType, kDimensionality>);
+    RadiusSearchResults results{};
+    radiusSearch(query, radius, sort, results);
+    return results;
+  }
+
+  /**
+   * @brief Returns the K closest points in the tree to a given query point.
+   *
+   * @param query The point to query the tree with.
+   * @param numClosest The number of closest points to return.
+   * @param results The K closest points, denoted by their index and squared distance. If the tree
+   * contains fewer than @param numClosest points, returns all the points in the tree instead.
+   */
+  template <typename EigenVectorType>
+  void knnSearch(EigenVectorType&& query, size_t numClosest, KnnSearchResults& results) const {
+    static_assert(detail::IsContiguousEigenVectorBaseWithLength<EigenVectorType, kDimensionality>);
+
+    numClosest = std::min(numClosest, getSize());
+    results.indices.resize(numClosest);
+    results.squaredDistances.resize(numClosest);
+    if (numClosest == 0) {
+      return; // Have to early-exit here because nanoflann doesn't support empty queries.
+    }
+
+    tree_.index->knnSearch(
+        query.data(), numClosest, results.indices.data(), results.squaredDistances.data());
+  }
+
+  /**
+   * @brief Returns the K closest points in the tree to a given query point.
+   *
+   * @param query The point to query the tree with.
+   * @param numClosest The number of closest points to return.
+   * @return The K closest points, denoted by their index and squared distance. If the tree contains
+   * fewer than @param numClosest points, returns all the points in the tree instead.
+   */
+  template <typename EigenVectorType>
+  KnnSearchResults knnSearch(EigenVectorType&& query, const size_t numClosest) const {
+    static_assert(detail::IsContiguousEigenVectorBaseWithLength<EigenVectorType, kDimensionality>);
+    KnnSearchResults results{};
+    knnSearch(query, numClosest, results);
+    return results;
+  }
+
+  /**
+   * @brief Returns the closest point in the tree to a given query point.
+
+   * @param query The point to query the tree with.
+   * @return The closest point. If the tree is empty, returns std::nullopt.
+   */
+  template <typename EigenVectorType>
+  std::optional<ClosestSearchResult> closestSearch(EigenVectorType&& query) const {
+    static_assert(detail::IsContiguousEigenVectorBaseWithLength<EigenVectorType, kDimensionality>);
+
+    if (isEmpty()) {
+      return std::nullopt;
+    }
+
+    ClosestSearchResult result{};
+    tree_.index->knnSearch(query.data(), /*num_closest=*/1, &result.index, &result.squaredDistance);
+    return result;
+  }
+
+ private:
+  EigenMatrixType points_;
+  TreeAdaptorType tree_;
+};
+
+template <typename ScalarType>
+using KdTree2 = KdTree<Eigen::Matrix<ScalarType, Eigen::Dynamic, 2>>;
+using KdTree2f = KdTree2<float>;
+using KdTree2d = KdTree2<double>;
+
+template <typename ScalarType>
+using KdTree3 = KdTree<Eigen::Matrix<ScalarType, Eigen::Dynamic, 3>>;
+using KdTree3f = KdTree3<float>;
+using KdTree3d = KdTree3<double>;
+
+} // namespace axel

include/axel/Log.h

@@ -0,0 +1,22 @@
+/*
+ * 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
+
+#if defined(AXEL_WITH_XR_LOGGER)
+
+#include <logging/Log.h>
+
+#else
+
+#define XR_LOGT(...)
+#define XR_LOGD(...)
+#define XR_LOGI(...)
+#define XR_LOGW(...)
+#define XR_LOGE(...)
+
+#endif

include/axel/MeshToSdf.h

@@ -0,0 +1,123 @@
+/*
+ * 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 <limits>
+#include <queue>
+#include <vector>
+
+#include <Eigen/Core>
+#include <gsl/span>
+
+#include "axel/BoundingBox.h"
+#include "axel/SignedDistanceField.h"
+#include "axel/TriBvh.h"
+
+namespace axel {
+
+/**
+ * Configuration parameters for mesh-to-SDF conversion.
+ */
+template <typename ScalarType>
+struct MeshToSdfConfig {
+  using Scalar = ScalarType;
+
+  /// Narrow band width around triangles (in voxel units)
+  Scalar narrowBandWidth = Scalar{1.5};
+
+  /// Maximum distance to compute (distances beyond this are clamped)
+  /// Set to 0 to disable clamping
+  Scalar maxDistance = Scalar{0};
+
+  /// Numerical tolerance for computations
+  Scalar tolerance = std::numeric_limits<Scalar>::epsilon() * Scalar{1000};
+};
+
+/**
+ * Convert a triangle mesh to a signed distance field using modern 3-step approach:
+ * 1. Narrow band initialization with exact triangle distances
+ * 2. Fast marching propagation using Eikonal equation
+ * 3. Sign determination using ray casting
+ *
+ * @param vertices Vertex positions as span (works with std::vector, arrays, subranges)
+ * @param triangles Triangle indices as span (indices must be valid within vertices)
+ * @param bounds Spatial bounds for the SDF
+ * @param resolution Grid resolution (nx, ny, nz)
+ * @param config Configuration parameters
+ * @return Generated signed distance field
+ */
+template <typename ScalarType>
+SignedDistanceField<ScalarType> meshToSdf(
+    gsl::span<const Eigen::Vector3<ScalarType>> vertices,
+    gsl::span<const Eigen::Vector3i> triangles,
+    const BoundingBox<ScalarType>& bounds,
+    const Eigen::Vector3<Index>& resolution,
+    const MeshToSdfConfig<ScalarType>& config = {});
+
+/**
+ * Convenience overload that computes bounds automatically from the mesh.
+ *
+ * @param vertices Vertex positions as span
+ * @param triangles Triangle indices as span
+ * @param resolution Grid resolution (nx, ny, nz)
+ * @param padding Extra space around mesh bounds (as fraction of bounding box size)
+ * @param config Configuration parameters
+ * @return Generated signed distance field
+ */
+template <typename ScalarType>
+SignedDistanceField<ScalarType> meshToSdf(
+    gsl::span<const Eigen::Vector3<ScalarType>> vertices,
+    gsl::span<const Eigen::Vector3i> triangles,
+    const Eigen::Vector3<Index>& resolution,
+    ScalarType padding = ScalarType{0.1},
+    const MeshToSdfConfig<ScalarType>& config = {});
+
+namespace detail {
+
+// ================================================================================================
+// FORWARD DECLARATIONS
+// ================================================================================================
+
+/**
+ * Initialize narrow band with exact triangle distances.
+ */
+template <typename ScalarType>
+void initializeNarrowBand(
+    gsl::span<const Eigen::Vector3<ScalarType>> vertices,
+    gsl::span<const Eigen::Vector3i> triangles,
+    SignedDistanceField<ScalarType>& sdf,
+    ScalarType bandWidth);
+
+/**
+ * Propagate distances from narrow band to entire grid using fast marching.
+ */
+template <typename ScalarType>
+void fastMarchingPropagate(SignedDistanceField<ScalarType>& sdf);
+
+/**
+ * Apply correct signs to distance field based on inside/outside classification.
+ */
+template <typename ScalarType>
+void applySignsToDistanceField(
+    SignedDistanceField<ScalarType>& sdf,
+    gsl::span<const Eigen::Vector3<ScalarType>> vertices,
+    gsl::span<const Eigen::Vector3i> triangles);
+
+/**
+ * Compute mesh bounding box from vertex spans.
+ */
+template <typename ScalarType>
+BoundingBox<ScalarType> computeMeshBounds(gsl::span<const Eigen::Vector3<ScalarType>> vertices);
+
+} // namespace detail
+
+// Type aliases for convenience
+using MeshToSdfConfigf = MeshToSdfConfig<float>;
+using MeshToSdfConfigd = MeshToSdfConfig<double>;
+
+} // namespace axel