Astronomical time series analysis with JAX

jaxoplanet is a functional-programming-forward implementation of many features from the exoplanet and starry packages built on top of JAX.

jaxoplanet includes fast and robust implementations of many exoplanet-specific operations, including solving Kepler’s equation, and computing limb-darkened light curves. Since jaxoplanet is built on top of JAX it has first-class support for hardware acceleration using GPUs and TPUs, and it also integrates seamlessly with modeling tools like NumPyro, and Flax.

For the most complete documentation, check out the documentation page at jax.exoplanet.codes.


You’ll first need to install JAX following the instructions in the JAX docs. For example, to install the CPU version of JAX, you can run:

python -m pip install "jax[cpu]"

Then install jaxoplanet with:

python -m pip install jaxoplanet

If you run into issues with installing jaxoplanet, take a look at the installation instructions.

Quick start#


While we don’t yet have a citation for jaxoplanet, please reference the GitHub repository if you find this code useful in your research. The BibTeX entry for the repo is:

  author       = {Soichiro Hattori and
                  Lionel Garcia and
                  Catriona Murray and
                  Jiayin Dong and
                  Shashank Dholakia and
                  David Degen and
                  Daniel Foreman-Mackey},
  title        = {{exoplanet-dev/jaxoplanet: Astronomical time series analysis with JAX}},
  month        = mar,
  year         = 2024,
  publisher    = {Zenodo},
  version      = {v0.0.2},
  doi          = {10.5281/zenodo.10736936},
  url          = {https://doi.org/10.5281/zenodo.10736936}


Copyright (c) 2021-2024 Simons Foundation, Inc.

jaxoplanet is free software made available under the MIT License. For details see the LICENSE file.

Table of contents#