Differentiable N-body code for Galactic Dynamics -- Odisseo

TL;DR

Constrained modeling of the Milky Way's gravitational potential using stellar streams requires differentiable, physics-based simulations. Odisseo provides a differentiable direct $N$-body solver in the JAX ecosystem with JIT, vectorization, and multi-GPU support, enabling gradient-based inference for Galactic dynamics. The work demonstrates near-linear multi-GPU scaling and two inference routes on a mock GD-1 stream—gradient-descent with bootstrapping and orbit fitting with Fisher contours—recovering host and progenitor properties. This establishes a scalable, end-to-end differentiable framework for Galactic dynamics that supports uncertainty quantification and rapid prototyping, with pathways to neural-network coupling and broader dynamical applications.

Abstract

We introduce \textsc{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. \textsc{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 a mock GD-1 stellar stream simulation, where we optimize four physical parameters via gradient descent: the accretion time and progenitor mass, as well as the masses of the host Navarro-Frenk-White (NFW) halo and Miyamoto-Nagai (MN) disk. \textsc{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.

Figures (5)

• Figure 1: Speedups for a single simulation distributed across multiple GPUs with different number of particles N. The dashed red line is the ideal linear scale, which is achieved for N > $10^4$.
• Figure 2: Inference pipeline. The goal is to generate $N$ bootstrap samples of the parameters. 1. From the observed stellar stream, we identify its Center of Mass (CoM). 2. Using an initial estimate $\hat{\Theta}_i = (\hat{t}_{\text{end}}, \hat{M}_{\text{Plummer}}, \hat{M}_{\text{NFW}}, \hat{M}_{\text{MN}})$, obtained through a grid search, we integrate the present-day CoM position backward in time as a single particle of mass $\hat{M}_{\text{Plummer}}$ for a duration $\hat{t}_{\text{end}}$, within an external potential parametrized by $\hat{M}_{\text{NFW}}, \hat{M}_{\text{MN}}$. 3. The recovered initial CoM position is populated with 1000 particles sampled from a Plummer sphere in equilibrium. 4. The resulting simulated stream is compared to the observed stream using the Maximum Mean Discrepancy (MMD), and the gradient of the loss is computed. 5. We perform 10 steps of vanilla Gradient Descent, repeating steps 2--4 internally. The parameters at the minimum MMD are taken as the final estimate for the sample $\Theta_i = (t_{\text{end}}, M_{\text{Plummer}}, M_{\text{NFW}}, M_{\text{MN}})$. 6. The observed stream is resampled according to observational errors to generate a new realization of the target, and the process is repeated from step 1.
• Figure 3: Bootstrapping inference results. The masses $M_{\text{MN}}$, $M_{\text{NFW}}$ and $M_{\text{Plummer}}$ are in $M_\odot$, while the total time of integration $t_{\text{end}}$ is in $Gyr$. The gray box represents the "prior" used to set up the initial grid search. The darker area in the diagonal plots represents one standard deviation from the maximum likelihood. The dashed lines indicate the true values.
• Figure 4: Fisher confidence contours around the Maximum Likelihood Estimate (MLE). All parameter values are shown in logarithmic (dex) units. The dashed lines are indicating the position of the true value.
• Figure 5: Snapshots at different times $t$ of the GD1 stream simulation. Milky Way's stellar disk is shown for scale and the center of mass (CoM) trajectory is shown as solid line.