Metadata-Version: 2.4
Name: jnkepler
Version: 0.2.10
Summary: Differentiable JAX N-body code for modeling nearly-Keplerian orbits
Author-email: Kento Masuda <kmasuda@ess.sci.osaka-u.ac.jp>
License: MIT
Project-URL: Homepage, https://github.com/kemasuda/jnkepler
Project-URL: Repository, https://github.com/kemasuda/jnkepler
Project-URL: Documentation, https://github.com/kemasuda/jnkepler
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax
Requires-Dist: numpyro
Requires-Dist: pandas
Requires-Dist: numpy
Requires-Dist: matplotlib
Requires-Dist: importlib_resources
Requires-Dist: corner
Requires-Dist: arviz
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: jaxoplanet; extra == "test"
Requires-Dist: setuptools<82; extra == "test"
Dynamic: license-file

# 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](https://jax.readthedocs.io/en/latest/index.html), 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](https://num.pyro.ai). 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 differentiable N-body model for analyzing transit timing variations (TTVs) and radial velocities (RVs) of multi-planet systems.
- **jnkepler.nbodytransit**: A differentiable photodynamical model. [`jaxoplanet`](https://jax.exoplanet.codes/en/latest/) needs to be installed for using this package.
- **jnkepler.nbodyrv**: A differentiable RV model taking into account mutual interactions between planets.

See [readthedocs](https://jnkepler.readthedocs.io/en/stable/) for more details.

## Installation

```pip install jnkepler```

### *CPU performance note*

`jnkepler` is primarily intended for CPU execution and supports recent JAX/jaxlib releases (tested through 0.12.0.1). With **both JAX and jaxlib >=0.12**, the default settings provide good CPU performance.

For older versions, use the matching `XLA_FLAGS` setting below to avoid an XLA CPU performance regression. Apply it **before importing JAX, jnkepler, NumPyro, or any other library that imports JAX**. In notebooks, restart the kernel first if these libraries have already been imported.

**JAX/jaxlib 0.7–0.11: increase the small-while-loop threshold**

```bash
export XLA_FLAGS="--xla_backend_extra_options=xla_cpu_small_while_loop_byte_threshold=65536"
```

Or in Python:

```python
import os
os.environ["XLA_FLAGS"] = (
    "--xla_backend_extra_options="
    "xla_cpu_small_while_loop_byte_threshold=65536"
)

import jax
import jnkepler
```

**JAX/jaxlib 0.4.32–0.6.x: use the legacy CPU runtime**

```bash
export XLA_FLAGS="--xla_cpu_use_thunk_runtime=false"
```

Or in Python:

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

import jax
import jnkepler
```

When upgrading from 0.6.x to 0.7–0.11, replace the legacy runtime flag with the threshold setting above. Once both JAX and jaxlib are >=0.12, remove these workaround settings.

See also [issue #30](https://github.com/kemasuda/jnkepler/issues/30) for jnkepler benchmarks and version comparisons, and [OpenXLA issue #50040](https://github.com/openxla/xla/issues/50040) for the upstream XLA CPU performance issue and fix.

### *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

Start with the step-by-step tutorials: [fit transit times](docs/quickstart.md), [write a NumPyro model](docs/numpyro_model.md), and [initialize, tune, and diagnose NUTS](docs/nuts.md).
For non-transiting planets, RVs, photodynamics, and pure Keplerian calculations, see [Models and extensions](docs/models.md).

The `examples/` directory contains the notebooks listed below. See the [example notebook guide](docs/examples.md) for their scope and requirements.

- **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)](https://ui.adsabs.harvard.edu/abs/2020AJ....159...57L/abstract)
- **TTV modeling** (Student's t likelihood):
  - `examples/v1298tau_ttv_student.ipynb`: reproducing the analysis in [Livingston et al. (2026)](https://www.nature.com/articles/s41586-025-09840-z)
  - `examples/trappist-1_ttv_student.ipynb`: reproducing the analysis in [Agol et al. (2021)](https://iopscience.iop.org/article/10.3847/PSJ/abd022)
  
- **Photodynamical modeling**: `examples/kep51_photodynamics_gp.ipynb`
  - SVI optimization & posterior sampling with NUTS
  - noise modeling using Gaussian Process with [`tinygp`](https://tinygp.readthedocs.io/en/stable/)

## Applications

- TOI-1136: TTV modeling of 6-planets in a resonance chain [[paper]](https://ui.adsabs.harvard.edu/abs/2022arXiv221009283D/abstract)
- TOI-2015: Joint TTV & RV modeling of a two-planet system [[paper]](https://arxiv.org/abs/2310.11775)
- Kepler-51: Four-planet modeling including JWST data [[paper]](https://arxiv.org/abs/2410.01625) [[repository]](https://github.com/kemasuda/kep51_jwst)
- K2-19: TTVs confirm 3:2 resonance [[paper]](https://arxiv.org/abs/2509.18031)
- TOI-4495: Photodynamical modeling of a pair of near-resonant sub-Neptunes [[paper]](https://arxiv.org/abs/2601.02665) [[repository]](https://github.com/kemasuda/toi4495-ttv)
- V1298 Tau: Four low-density planets transiting a young star [[paper]](https://www.nature.com/articles/s41586-025-09840-z) [[repository]](https://github.com/kemasuda/v1298tau-ttv); see also `examples/v1298tau_ttv_student.ipynb` in this repo
- TOI-2076: TTV modeling of a young, near-resonant planetary system [[paper]](https://ui.adsabs.harvard.edu/abs/2026arXiv260302550W/abstract)
- TOI-791: TTV masses of two extremely low-density, Jupiter-sized planets [[paper]](https://arxiv.org/abs/2606.30016)
- TOI-3686: Photodynamical and RV modeling of a hot Jupiter and sub-Neptune in a 2:1 mean-motion resonance [[paper]](https://arxiv.org/abs/2610.28533)

## References

If you use `jnkepler` in your research, please consider citing the references below.

- Masuda et al. (2024), [A Fourth Planet in the Kepler-51 System Revealed by Transit Timing Variations](https://ui.adsabs.harvard.edu/abs/2024AJ....168..294M/abstract), AJ 168, 294
- Masuda (2025), [jnkepler: Differentiable N-body model for multi-planet systems](https://ui.adsabs.harvard.edu/abs/2025ascl.soft05006M/abstract),  Astrophysics Source Code Library, ascl:2505.006.
