Search papers, labs, and topics across Lattice.
This paper introduces jaxsgp4, a JAX-based reimplementation of the SGP4 orbital propagation algorithm, optimized for GPU execution. By refactoring SGP4 into a pure functional paradigm, the authors leverage JAX's JIT compilation and automatic vectorization to achieve massive parallelism. The results demonstrate a 1500x speedup over traditional C++ baselines, propagating the entire Starlink constellation to 1,000 future time steps in under 4ms on a single A100 GPU, while also demonstrating the viability of 32-bit precision for further throughput gains.
Propagating mega-constellations is now 1500x faster thanks to a JAX-based SGP4 reimplementation, making large-scale collision avoidance tractable.
As the population of anthropogenic space objects transitions from sparse clusters to mega-constellations exceeding 100,000 satellites, traditional orbital propagation techniques face a critical bottleneck. Standard CPU-bound implementations of the Simplified General Perturbations 4 (SGP4) algorithm are less well suited to handle the requisite scale of collision avoidance and Space Situational Awareness (SSA) tasks. This paper introduces \texttt{jaxsgp4}, an open-source high-performance reimplementation of SGP4 utilising the \texttt{JAX} library. \texttt{JAX} has gained traction in the landscape of computational research, offering an easy mechanism for Just-In-Time (JIT) compilation, automatic vectorisation and automatic optimisation of code for CPU, GPU and TPU hardware modalities. By refactoring the algorithm into a pure functional paradigm, we leverage these transformations to execute massively parallel propagations on modern GPUs. We demonstrate that \texttt{jaxsgp4} can propagate the entire Starlink constellation (9,341 satellites) each to 1,000 future time steps in under 4 ms on a single A100 GPU, representing a speedup of $1500\times$ over traditional C++ baselines. Furthermore, we argue that the use of 32-bit precision for SGP4 propagation tasks offers a principled trade-off, sacrificing negligible precision loss for a substantial gain in throughput on hardware accelerators.