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 jaxoplanetIf you run into issues with installing jaxoplanet, take a look at the installation instructions.
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:
@software{jaxoplanet,
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.