jaxsgp4: GPU-accelerated mega-constellation propagation with batch parallelism

TL;DR

The paper presents jaxsgp4, a GPU-accelerated, highly parallel implementation of the SGP4 orbital propagation algorithm that significantly outperforms traditional methods, enabling rapid large-constellation simulations.

Contribution

It introduces jaxsgp4, a novel JAX-based, GPU-optimized SGP4 implementation that achieves massive speedups and supports large-scale satellite constellation propagation.

Findings

Propagates 9,341 Starlink satellites in under 4 ms on a single GPU.

Achieves a 1500x speedup over traditional C++ implementations.

Uses 32-bit precision to balance accuracy and performance.

Abstract

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.