exoplanet-dev / jaxoplanet

Astronomical time series analysis with JAX
https://jax.exoplanet.codes
MIT License
41 stars 12 forks source link
astronomy astrophysics bayesian-inference exoplanets jax python time-series

jaxoplanet

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.

Installation

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, specifically on ARM (i.e., M series chips) Macs, take a look at the installation instructions.

:class: tip

💽 After [installing](doc:install) `jaxoplanet`, head over to the [Quickstart](quickstart) page to see some of its features.

🚩 If you're running into some problems with `jaxoplanet`, check out the [common issues](doc:commonissues) page for some general tips and tricks.

⚡ In the [Tutorials](tutorials) section we've added a few tutorials showcasing some common astronomy problems where `jaxoplanet` might come in handy!

🖥️ Take a look [here](doc:api) to see a user-friendly API reference for commonly used objects. The full API reference is also available!

Attribution

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}
}

License

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

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