Differentiable N-body code for Galactic Dynamics - Odisseo

Giuseppe Viterbo, Tobias Buck

Keywords: Differentiable simulator, Galactic Dynamics, N-body simulator, Gradient Based Optimization

TL;DR: We present Odisseo, a differentiable N-body code built in JAX, enabling scalable simulations and parameter estimation in Galactic Dynamics.

Abstract: We introduce Odisseo (Optimized Differentiable Integrator for Stellar Systems Evolution of Orbits), a differentiable N-body code designed to constrain the gravitational potential of the Milky Way (MW) through dynamical modeling of accreted structures such as stellar streams. Odisseo is implemented in JAX, enabling just-in-time compilation, automatic differentiation, and hardware acceleration on GPUs and TPUs. The code features efficient, fully vectorized force calculations and exhibits near-linear scaling when distributing a single simulation across multiple GPUs, making it suitable for large-scale optimization tasks. As a demonstration, we present a case study using the GD-1 stellar stream, where we optimize four physical parameters via gradient descent: the accretion time and progenitor mass, as well as the masses of the host galaxy’s Navarro-Frenk-White (NFW) halo and Miyamoto–Nagai (MN) disk. Odisseo accurately recovers stream morphology and underlying parameters in a differentiable and scalable framework, providing a powerful tool for dynamical studies of the Milky Way and its accreted substructures. The code is available on GitHub: https://github.com/vepe99/Odisseo.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 6