search

exoplanet-dev / jaxoplanet

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

Vectorize light curve functions #141

Closed dfm closed 9 months ago

dfm commented 9 months ago

This PR updates the interface introduced by #129 to remove the requirement that light curves always be manually vmapped. But, it still allows the logic within the light curve computation to assume that it only needs to operate on a scalar time. I think this is the best of both worlds because it makes the implementation of light curves a little more straightforward while keeping a nice user interface that can operate on arrays of times.

To do this, I added a light_curves.utils.vectorize decorator which handles the vmapping and units. We can't use jnp.vectorize (which would have been my preference) because we need to support Quantity input and output.

lgrcia commented 9 months ago

Looks nice! Once merged I'll try to use that for the starry light curves functions too.

dfm commented 9 months ago

I'm going to merge this for now so that we can keep rolling with some other PRs, but I'd love to hear your thoughts later when you have a chance @soichiro-hattori!

soichiro-hattori commented 9 months ago

@dfm: This all looks great!