Skip to content

PolytopAX: GPU-accelerated convex polytope approximation and geometry library powered by JAX

License

Notifications You must be signed in to change notification settings

lv416e/polytopax

πŸ”· PolytopAX

GPU-accelerated differentiable convex hull computation powered by JAX

License PyPI Tests Lint Docs Release

PolytopAX brings modern computational geometry to the JAX ecosystem, enabling differentiable convex hull computation with GPU acceleration and automatic differentiation for machine learning applications.

✨ Key Features

  • πŸš€ GPU-Accelerated: Leverage JAX/XLA for high-performance computation
  • πŸ”„ Differentiable: Full compatibility with JAX transformations (jit, grad, vmap)
  • 🎯 ML-Ready: Seamless integration into machine learning pipelines
  • πŸ“¦ Easy to Use: Both functional and object-oriented APIs
  • πŸ”¬ Research-Grade: Built for computational geometry research and applications

πŸš€ Quick Start

Installation

pip install polytopax

Basic Usage

import jax.numpy as jnp
import polytopax as ptx

# Generate random points
points = jax.random.normal(jax.random.PRNGKey(0), (100, 3))

# Compute convex hull
hull_vertices = ptx.convex_hull(points)
print(f"Hull has {len(hull_vertices)} vertices")

# Object-oriented API
hull = ptx.ConvexHull.from_points(points)
print(f"Volume: {hull.volume():.4f}")
print(f"Surface area: {hull.surface_area():.4f}")

# Check point containment
test_point = jnp.array([0.0, 0.0, 0.0])
is_inside = hull.contains(test_point)

Differentiable Optimization

import jax

# Differentiable volume computation
def volume_loss(points):
    hull = ptx.ConvexHull.from_points(points)
    return -hull.volume()  # Maximize volume

# Compute gradients
grad_fn = jax.grad(volume_loss)
gradients = grad_fn(points)

# Batch processing with vmap
batch_points = jnp.stack([points, points + 0.1])
batch_volumes = jax.vmap(lambda p: ptx.ConvexHull.from_points(p).volume())(batch_points)

πŸ”§ API Overview

Core Functions

  • convex_hull(points) - Compute convex hull vertices
  • point_in_convex_hull(point, vertices) - Point containment test
  • convex_hull_volume(vertices) - Volume computation
  • convex_hull_surface_area(vertices) - Surface area computation

ConvexHull Class

  • ConvexHull.from_points(points) - Create from point cloud
  • .volume() - Compute volume
  • .surface_area() - Compute surface area
  • .contains(point) - Test point containment
  • .centroid() - Compute centroid

🎯 Use Cases

  • Robotics: Path planning and obstacle avoidance
  • Machine Learning: Constraint optimization and geometric deep learning
  • Computer Graphics: Collision detection and rendering
  • Computational Physics: Molecular dynamics and simulations
  • Finance: Risk management and portfolio optimization

πŸ—οΈ Why PolytopAX?

Feature PolytopAX SciPy Qhull
GPU Acceleration βœ… ❌ ❌
Differentiable βœ… ❌ ❌
JAX Integration βœ… ❌ ❌
Batch Processing βœ… ❌ ❌
ML Pipeline Ready βœ… ❌ ❌

πŸ“– Documentation

πŸ”¬ Examples

Explore comprehensive examples in the examples/ directory:

πŸ›£οΈ Roadmap

  • v0.1.0 βœ… Core differentiable convex hull algorithms
  • v0.5.0 πŸ”„ Exact algorithms (Quickhull, incremental)
  • v0.8.0 πŸ“‹ Advanced operations (intersections, Minkowski sums)
  • v1.0.0 🎯 Riemannian manifold integration (β†’ GeomAX)

🀝 Contributing

We welcome contributions! Please see our Contributing Guidelines and Development Setup.

πŸ“„ License

Licensed under the Apache License, Version 2.0. See LICENSE for details.

πŸ”— Links


Built with ❀️ for the JAX and computational geometry communities

About

PolytopAX: GPU-accelerated convex polytope approximation and geometry library powered by JAX

Resources

License

Code of conduct

Security policy

Stars

Watchers

Forks

Packages

No packages published

Contributors 2

  •  
  •  

Languages