jnkepler 0.2.7

Differentiable JAX N-body code for modeling nearly-Keplerian orbits

jnkepler

A differentiable N-body model for multi-planet systems.

jnkepler is a Python package for modeling photometric and radial velocity data of multi-planet systems via N-body integration. Built with JAX, it leverages automatic differentiation for efficient computation of model gradients. This enables seamless integration with gradient-based optimizers and Hamiltonian Monte Carlo methods, including the No-U-Turn Sampler (NUTS) in NumPyro. The package is particularly suited for efficiently sampling from multi-planet posteriors involving a large number of parameters and strong degeneracy.

Subpackages

• jnkepler.jaxttv: A differentialble N-body model for analyzing transit timing variations (TTVs) and radial velocities (RVs) of multi-planet systems.
• jnkepler.nbodytransit: A differentialble photodynamical model. jaxoplanet needs to be installed for using this package.
• jnkepler.nbodyrv: A differentiable RV model taking into account mutual interactions between planets.

See readthedocs for more details.

Installation

pip install jnkepler

CPU performance note

If you use jnkepler on CPU with JAX ≥ 0.4.32, the default thunk runtime in the CPU backend can make computations much slower, especially when computing gradients.

To avoid this, disable the thunk runtime by setting the following environment variable before importing jax:

export XLA_FLAGS="--xla_cpu_use_thunk_runtime=false"

Or inside Python:

import os
os.environ["XLA_FLAGS"] = "--xla_cpu_use_thunk_runtime=false"
import jax

If this is not done, jnkepler will issue a warning on import. Please note that this workaround is intended for JAX < 0.7; this is why jnkepler currently requires jax<0.7.

Note on the transit-finding algorithm (since v0.2.5)

In JaxTTV.get_transit_times_obs, the default transit-time solver is now "fast". This method is optimized for efficient evaluation near the observed transit times, which is typically sufficient for inference because models far from the data have very low likelihood. For models that are substantially offset from the observed transit times, however, the fast method can be less accurate than the Newton-based method. Use transit_time_method="newton" when robust accuracy is needed in that regime.

For this reason, when using the fast method in NUTS, initialize the sampler from a good solution, such as a least-squares optimum as in the example notebooks, rather than from a diffuse or poor initial guess. If multimodality is a major concern, consider a preliminary grid/global search before running NUTS, rather than relying on NUTS alone to discover all modes.

Examples

Explore example notebooks in the examples/ directory to see jnkepler in action:

• minimal example: examples/minimal_example.ipynb

  • computing transit times and RVs
  • plotting TTVs
  • adding a non-transiting planet

• TTV modeling (normal likelihood): examples/kep51_ttv_normal.ipynb

  • posterior sampling with NUTS
  • reproducing the result in Libby-Roberts et al. (2020)

• TTV modeling (Student's t likelihood):

  • examples/v1298tau_ttv_student.ipynb: reproducing the analysis in Livingston et al. (2026)
  • examples/trappist-1_ttv_student.ipynb: reproducing the analysis in Agol et al. (2021)

• Photodynamical modeling: examples/kep51_photodynamics_gp.ipynb

  • SVI optimization & posterior sampling with NUTS
  • noise modeling using Gaussian Process with tinygp

Applications

• TOI-1136: TTV modeling of 6-planets in a resonance chain [paper]
• TOI-2015: Joint TTV & RV modeling of a two-planet system [paper]
• Kepler-51: Four-planet modeling including JWST data [paper] [repository]
• K2-19: TTVs confirm 3:2 resonance [paper]
• TOI-4495: Photodynamical modeling of a pair of near-resonant sub-Neptunes [paper] [repository]
• V1298 Tau: Four low-density planets transiting a young star [paper] [repository]; see also examples/v1298tau_ttv_student.ipynb in this repo
• TOI-2076: [paper]

References

• Masuda et al. (2024), A Fourth Planet in the Kepler-51 System Revealed by Transit Timing Variations, AJ 168, 294
• Masuda (2025), jnkepler: Differentiable N-body model for multi-planet systems, Astrophysics Source Code Library, ascl:2505.006.