Implement semi-analytic root solver in Zhang (2022)

[kmzzhang]

Hi Frac,

As we discussed, the semi-analytic root solver may substantially faster than the AE solver. Let's see if we can implement it in JAX and see how it works out.

https://github.com/kmzzhang/analytic-lensing
https://arxiv.org/abs/2207.12412

Keming

[fbartolic]

Hi @kmzzhang ,

If I got it right, the algorithm is:

1. Solve eq. for 2 single lens images, pick image with Re(z) < 0.
2. Intialize Newton or Laguerre solver with that value and find a single isolated root of the 5th order complex polynomial.
3. Get coefficients for the 4th order complex polynomial using polynomial division.
4. Solve analytically for the 4 roots and combine with the first root to obtain all images.

This should be trivial to do in JAX to integrate within the caustics API. I'll work it out in a notebook and post here when I obtain some results on performance improvements.

Is this method not what Skowron & Gould are doing in their optimized zroots routine?

[kmzzhang]

Hi @fbartolic,

The algorithm is exactly as you described. As for step 3, the quartic coefficients can be pre-determined by substituting z with z+z0 in the quintic polynomial (z0 is found in step 2) and dropping the zeroth-order coefficient that will be close to zero. See https://github.com/kmzzhang/analytic-lensing/blob/main/exact.py#L144.

Nope, Skowron & Gould (neither do other algorithms I know of) does not use the single lens image location as an initial guess. If I understand correctly, it uses a combination of Newton + Laguerre's method for point source evaluations. When there exists approximate roots for finite source calculations from previous locations, It numerically solves for the two "in caustic" images, factors them out, and analytical solves for the three other images.

[fbartolic]

Here's a first attempt:
https://gist.github.com/fbartolic/f1021daf7dd44021b3befce763db6819

I'm getting the wrong answer for the 4 roots, I'm guessing because we define the coordinate system differently. In caustics the coordinate system is set up such that the origin is at the midpoint of the two lenses which are both on the real line located at a (lens with mass fraction e1) and -a (lens with mass fraction 1-e1), where s=2a. I'll get back to this some time tomorrow.

[kmzzhang]

I'm working on correcting it and will upload soon

[kmzzhang]

Ok I fixed it: https://gist.github.com/kmzzhang/831bbc3282cb0c08ba702631e2023842
Changes are:

1. coordinate transformation: also has to scale by (1-e1)**0.5 because units are defined in Einstein radius of primary mass and total mass.
2. I see you used the original "NKrvavica/fqs" code --- it doesn't work out of the box for complex polynomials. Specifically its cubic solver that I had to rewrite.

[fbartolic]

coordinate transformation: also has to scale by (1-e1)**0.5 because units are defined in Einstein radius of primary mass and total mass.

Thanks for spotting this!

For consistency, we should derive the coefficients of the reduced polynomial using SymPy like I did here.

I see you used the original "NKrvavica/fqs" code --- it doesn't work out of the box for complex polynomials. Specifically its cubic solver that I had to rewrite.

Oh i see. I just copied that code blindly because it was quicker to port it to JAX. We don't need to vectorize the functions
at this level because it is better to use vmap to vectorize over the coefficients (at least on the CPU).

Have a look at the the revised notebook. Initial test looks good, semi-analytic method is accurate to within 1e-10 with fixed nr. of 15 Newton iterations. Speedup over Aberth-Ehrlich is less impressive than I though though, 0.67 ms vs 1.05ms for evaluating the roots of 200 polynomials very close to the caustics. This is without custom initialization and sequential updating.

I think it's worth the effort of implementing this properly but it doesn't seem like we'd get a speedup above a factor of 2 or so.

It's possible that I made a mistake somewhere though. I'll do a proper comparison tomorrow night.

[kmzzhang]

Good to hear! I'm sure it can be further optimized since it's currently 2x slower than my python implementation on my machine. I've been deriving coefficients in Mathematica -- I can replace it later with your parameterization.

I've created a draft pull request #17 for a shared platform for us to work together on. You should have push access to that branch. Let me know if you have other preferences.