Zernike enhancements

[jmeyers314]

Two tweaks to the Zernike library:

• Adds a DoubleZernike.xycoef() method to directly get the pupil coefficient array at a given field position of a double zernike, without going through a galsim.Zernike object.
• Enables large jmax for annular Zernikes. Addresses #1326 .

I found a closed-form for a certain subset of annular Zernikes where the radial and azimuthal orders are equal, so used that as a test of annular Zernikes Z66, Z231, Z861, Z1891, and Z3321. GalSim's annular Zernike library reproduces these to tolerances of part per 10^12, 10^12, 10^9, 10^6, and 10^3 respectively. I also tried Z5151. This value of j is apparently well beyond the accuracy of the library though and yields errors in the ~factor of 10 range.

These were all for R_inner/R_outer = 0.5; I'd guess the accuracies probably improve as this ratio goes to 0 (annular Zernikes go to circular Zernikes), but didn't explicitly check.

[rmjarvis]

Somewhat confusing reuse of variable n. Probably should rename the other one. (k maybe?)

[rmjarvis]

Suggested change

	print(f"Z{j} => (n, m) = ({n}, {n})")
	print(f"Z{j} => (n, m) = ({n}, {m})")

Probably followed by assert n == abs(m) since AFAIU, that is the intent of these choices for j.

[rmjarvis]

I don't think we expect this to converge particularly well when r>1, right? So should probably limit this to have x**2 + y**2 < 1.
When I tried it, I was able to use significantly lower tolerances:

(10, 1.e-12)
(20, 1.e-12)
(40, 1.e-10)
(60, 1.e-8)
(80, 1.e-6)
(100, 2.e-3)

So it is still becoming less accurate as j get very large. But not as catastrophically bad as this test implies.

[rmjarvis]

The ones that are least accurate are always ones with r close to 1. Like 0.95 or so. This is probably a clue about where the recursion is going unstable. I might play around with the recursion formula to see if I can harden it up at all.

[jmeyers314]

Yeah, limiting to 0.5 < r < 1.0 might even be appropriate. Though, OTOH, these are just 2d polynomials. They're orthogonal over an annulus but defined everywhere.

The particular polynomials here are proportional to r^n, so that might explain why the large r are trickier. Other values of j will presumably produce polynomials with larger amplitudes near r=0.

In case it's helpful while you're looking at recursions, here's the closed-form table I got out of Mathematica up to Z28:

[rmjarvis]

I think the problem is the conversion of the coef array from rrsq format to xy format. The xy version has lots of very large coefficients with alternating sign. The native rho/rhosq version is much more stable numerically.

If I change evalCartesian to:

        rho = (x + 1j*y) / self.R_outer
        rhosq = np.abs(rho)**2
        ar = _noll_coef_array(len(self.coef)-1, self.R_inner/self.R_outer).dot(self.coef[1:])
        return horner2d(rhosq, rho, ar).real

then all the tests pass at 1.e-12 tolerance, except Z5151, which needed 1.e-11. (It's also massively faster than the current implementation -- I guess the rrsq_to_xy function must be a slow step?)

I don't know if this means we should actually make the above change though. There might be other downsides (e.g. not getting to use the C++ layer horner2d implementation, at least without a little more work). But I suppose we could at least provide another method that would evaluate this way for users who want more accuracy at very high Zernike order.

[jmeyers314]

Interesting. I sped up rrsq_to_xy significantly by rewriting it as a cached recursive function. Now the tradeoff seems to be roughly:

The xy array is about 2x the effort to compute as the rrsq array, which makes sense since the xy is derived from the rrsq. However, most of this effort is cached so I don't think there's a significant difference after the initial Zernike constructions.

The xy approach (so c++ horner) is about 10x faster when there are ~100 evaluation points. However, on my machine (M3 Pro) the complex-python approach actually runs faster when the number of evaluation points is ~10_000 or more (!). Maybe this is because the rrsq array is ~2x smaller than the xy array?

Given the improved accuracy and the above, I tentatively vote make the rrsq approach the default. I want to benchmark some donut fitting first though.

[jmeyers314]

Hmmm.... Maybe the horner step depends on jmax too. I just ran some donut image forward models (jmax ~65 and npoints ~10000) and they were ~4x faster using the xy approach. So maybe need to leave that as an expert option.

[rmjarvis]

Either way is fine with me. As long as there is a public API way to get the calculation that is accurate for high order.

I think most use cases won't care about the accuracy as much as the speed. But it seems important to have a way to get the accurate one.