cs.DCMar 29, 2026

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

Authors: Charlotte Priestley, Will Handley

Organizations: Kavli Institute for Cosmology, University of Cambridge, Cambridge, UK · Institute of Astronomy, University of Cambridge, Cambridge, UK · PolyChord Ltd., London, UK

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 jaxsgp4, an open-source high-performance reimplementation of SGP4 utilising the JAX library. 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 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×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.

Figures & tables

Figure 1: Scaling performance comparison between JAX /GPU and C++/CPU implementations of the SGP4 algorithm, for a single satellite propagated to multiple times (left) and multiple satellites propagated to a single time (right). jaxsgp4 exploits the potential of modern GPUs to perform massively-parallel computations, exhibiting a flat scaling regime where increased workload does not increase wall-clock time until hardware saturation (top). The ‘break-even’ point for each GPU (Speedup=1, denoted by dashed line in bottom panels), at which the benefit of parallel computation for large batch sizes overcomes the initial dispatch overhead, occurs at batch sizes of ∼\sim 300-500. In the bottom panels, we also specify the maximum speedup achieved by jaxsgp4 over the C++ baseline. This occurs in the linear hardware saturation regime, at which point the GPUs are operating at maximum computational efficiency.
Figure 2: GPU-accelerated JAX vs C++ SGP4 performance comparison. Blue indicates JAX is faster, Red indicates C++ is faster. The standard C++ implementation is inherently single-threaded and was run on a standard CPU while the JAX benchmark was run on an NVIDIA A100 GPU. Log-spaced axes show the relative performance of each implementation across different scales.
Figure 3: Relative error accumulation over two weeks running jaxsgp4 with 32 bit precision versus the standard SGP4 implementation which uses 64 bit precision. The median and 5 th - 95 th percentile regions are highlighted in blue. The dashed line depicts a conservative lower bound on physical error growth due to limitations of the SGP4 model itself. jaxsgp4 errors are order 1e-9, 1e-12 for position and velocity respectively (floating point) when the same analysis is performed using 64 bit precision.
More in this topic

GPU Acceleration

GPU acceleration uses parallel processors, specialized kernels, and memory-aware systems to speed computation across AI and scientific workloads.

See all 169 papers

Explore similar work

CardsList
Jul 13, 2026

WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled LDL⊤LDL^\top Factorization

This paper introduces numerical optimizations for maximizing throughput on GPU when solving large batches (10,000 to over 100,000) of sequential quadratic programming (SQP) iterations, where all problems have the same structure. The optimizations are implemented in a toolbox WarpMPC for model-predictive control (MPC) in JAX and Warp. Based on the insight that all MPC problem instances in a batch share the same sparsity in time, cost, and constraints, we propose unrolling sparse linear factorizations and solves, which dominate alternating direction method of multipliers (ADMM) solver runtime. We avoid memory access bottlenecks and wasting computations via optimized memory layout, padding-reducing segmentation of the unrolled factorization, and dependency level scheduled backsolves, additionally accelerating sensitivity computation. We achieve throughputs of 8,000 to 250,000 SQP iterations per second on nonlinear cartpole, quadrotor, and humanoid robot benchmarks, outperforming baselines by 3×\times to 25×\times. We illustrate practical usefulness by synthesizing a dataset and training a neural network approximation of an MPC in under 4 minutes that stabilizes a nano quadrotor in hardware experiments.

Henrik Hose, Se Hwan Jeon, Charles Khazoom +2

Institute for Data Science in Mechanical Engineering (DSME), RWTH Aachen University, Germany · Department of Mechanical Engineering, Massachusetts Institute of Technology, Cambridge MA, USA.

GPU Kernel OptimizationGPU AccelerationNonlinear MPCAlternating Direction Method of MultipliersQuadratic Programming
AskPDF
Jun 4, 2026

Accelerated Fourier SAT (AFSAT): Fully Realising a GPU-based Symmetric Pseudo-Boolean SAT Solver

We present Accelerated Fourier SAT (AFSAT), a GPU-accelerated solver for pseudo-Boolean satisfiability based on continuous local search (CLS). AFSAT realises the proof-of-concept approach, FastFourierSAT, into a fully-engineered solver supporting any heterogeneous mixture of symmetric constraint types and lengths within a single problem instance. Using the JAX compiler, AFSAT leverages pure function composition, automatic vectorisation, automatic differentiation, and just-in-time (JIT) compilation to perform massively parallel CLS across batches of candidate assignments. We demonstrate substantially improved numerical stability, runtime performance, and memory efficiency over the proof-of-concept. We achieve this by way of identifying and addressing various limitations that arise from memory latency and floating-point representation, as well as leveraging automatic parallelisation and compact representations. The inherent representational and stability limitations of floating point are partially addressed by a tailored discrete Fourier transform implementation. We achieve near-linear throughput when scaling to multiple accelerators via JAX array sharding.

Cody J Christopher, Charles Gretton

School of Computing, Australian National University, Canberra, Australia.

GPU AccelerationHigh-Performance ComputingConstraint Satisfaction
AskPDF
May 31, 2026

Time-Optimal Collision Avoidance Via a Greedy Polynomial Backward Sweep

Spacecraft collision avoidance for low-thrust satellites often requires determining not only how to maneuver, but also how late a maneuver can begin while still ensuring safety. This paper presents a greedy time-optimal (GTO) backward-sweep method to find the latest maneuver initiation time. The method starts from the nominal time of closest approach and iteratively propagates the maneuver backward in time, selecting at each step the thrust direction that locally minimizes the chosen danger metric. Differential algebra is used to efficiently propagate state sensitivities and update the time of closest approach online. The method is tested on a large dataset of conjunctions, using both miss distance and probability of collision as safety metrics. The approach achieves accurate results and only a small loss of optimality relative to an optimal-control benchmark, while retaining runtimes suitable for on-board implementation.

Zeno Pavanello, Frank De Veld, Roberto Armellin

Department of Aerospace Science and Technology (DAER), Politecnico di Milano, via Privata Giuseppe La Masa 34, Milano, 20156, Italy · Netherlands Aerospace Centre (NLR), Anthony Fokkerweg 2, Amsterdam, 1059 CM, The Netherlands · Te Pūnaha ¯Atea Auckland - Space Institute (TPA-SI), The University of Auckland, 20 Symonds Street, Auckland, 1010, New Zealand

Collision AvoidanceTrajectory OptimizationOptimal Control
AskPDF