exoplanet-dev / jaxoplanet

Astronomical time series analysis with JAX
https://jax.exoplanet.codes
MIT License
38 stars 12 forks source link

API design #44

Closed lgrcia closed 1 month ago

lgrcia commented 1 year ago

I am opening this thread to discuss about the design of jaxoplanet's API. Do we want to merge the exoplanet and starry APIs in a new one for jaxoplanet, or do we want to keep both well separated?

I think exoplanet only deals with limb-darkened light-curves while starry allows for more things (e.g. more complex maps, reflected ones... etc.). Because the two rely on different implementations (same methods but lighter formulation for exoplanet thanks to Agol et. al 2020), I don't know if copying exoplanet's API will translate well in making starry-like features available in jaxoplanet.

Ideally, I think jaxoplanet must have a (different) unified API, with sensitive defaults for the light curves methods depending on the map of the central body (e.g. if only limb-darkening needs to be modeled). Of course, that requires a bit of API rethinking. Curious to hear your thoughts about it.

dfm commented 1 year ago

Thanks for starting this! My goal is to support at least a significant subset of the starry functionality directly within jaxoplanet rather than as a separate package. For example, right now there is substantial overlap between:

In the short term, @soichiro-hattori and I are planning on removing quad.py since limb_dark.py should have identical performance for quadratic models. In the longer term, it's possible that we could also deprecate limb_dark.py in favor of starry.py, but we'd need to think about performance.

As for the front end, I agree that we'll want to think about how to best expose the full interface, but I'm actually not to worried because some of the design issued faced in starry won't be so important here, since we can make a lot of decisions at JIT compilation time.

lgrcia commented 1 year ago

Thanks for your answer @dfm! What starry design issues are you thinking about? (looking for a starting point to improve the front-end later).

Needless to say I am fully in favor of a common package! :)

dfm commented 1 year ago

For performance reasons, starry needed to include many different compiled "ops". For example, limb darkening needed to be implemented separately from the general starry matrices. In this case, I think we can be a bit more flexible without having a proliferation of back-end models. Furthermore, things like wavelength-dependent maps can be efficiently implemented using jax.vmap rather than yet another custom op! I think that this flexibility means that we have more options for how to define the front-end, but it doesn't necessarily mean that we'll do it right :D