exoplanet-dev / jaxoplanet

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

Optimizing starry impl #204

Open dfm opened 1 month ago

dfm commented 1 month ago

These optimizations start to narrow the gap between the limb_dark and starry solution vector computational performance. The main differences:

  1. Using the closed form solution for s0 instead of the numerical one.
  2. In the previous implementation a lot of zeros were being summed in P integral for the diagonal case. The key is to just never compute the off diagonal args, and only update the indices accordingly.

Still to do:

  1. We should be able to use the closed form s2 as well, but it wasn't clear to me where it would go.
  2. Figure out the shapes. I don't totally understand the logic used for determining the "diagonal" indices, but why do we get a 4 element solution vector for order=2, instead of 3 elements like in the limb_dark case?
lgrcia commented 1 month ago

Thanks so much for having a look at that @dfm!

Your two questions are linked. Agol basis is

$$(1\;z\;z^2\;z^3\;\dots\;z^n)^T\text{,}$$

and Luger's basis is

$$(1\;x\;z\;y\;x^2\;xz\;xy\;yz\;y^2\dots)^T\text{.}$$

Luger's basis does not contain powers of $z$ greater than 1. So to convert vectors from Agol's basis to Luger's basis we have to use the fact that, on the unit sphere, $z^2 = 1 - x^2 - y^2$. For example, a surface described in Agol's basis

$$(1\;z\;z^2)^T$$

can be minimally described in Luger's basis by

$$(1\;z\;x^2\;y^2)^T\text{.}$$

This explains the difference in shapes between the two basis vectors.

Edit: in #203, I use that to find the transformation from one basis to the other, so that Agol's solution vector can be directly used in the starry version of the light curve.

soichiro-hattori commented 3 weeks ago

Hey @dfm and @lgrcia,

I've been trying to get a better understanding of some of the proposed changes in this PR and #203, and how it relates to the failed tests.

  1. Am I correct that the change in the for loop lower limit from 0 to 1 in both the p_integral and q_integral functions come from using the closed-form solution for the s0 component?
  2. The current code fails because of the shape mismatch for the sT @ A2 operation in the surface_light_curve function. My guess is that it's due to how sT is now calculated using Agol's basis while A2 is still in the Luger basis? @lgrcia: Do the newer commits in #203 fix this issue with another transformation?
lgrcia commented 3 weeks ago

@soichiro-hattori, I think this PR was more like a draft and indeed, it is missing some pieces. After fixing the bugs I reported most of the findings and changes back to #203.