Basis in the Long Term

[Jordan-Dennis]

Hi all,
We have been discussing this in the slack but I thought that I would add this discussion to elucidate further.

Short term

Basis implementation that can translate, magnify and (not yet) rotate. This exists minus the rotation and is probably what we will need to resort to using for the short term science targets. The rotation should just be multiplying the (2, npix, npix) array of coordinates by a (2, 2) standard rotation matrix. Having transformed the coordinate system we can repeat the calculation for the the aperture. This will work for any linear transformation of space because these are isomorphic to (2, 2) matrices. The reason I see this as a short term attempt is because it cannot cope with non-linear transformation.

Long term

The plan is to internally parameterise the aperture by it's vertices. We are not taking gradients or optimising with respect to the vertices. I decided to use this implementation because it is significantly easier to perform non-linear transformations on the vertices and it automatically copes with arbitrary apertures much more easily.

If we want to allow for transformations that change the "radius" parameter we need to be able to magnify the Zernikes. By radius I mean the radius of the smallest circle that can completely enclose the aperture. I do not think that an approximation avoiding this is going to be valid, even in the small regime.

There is also the question of are non-linear transformations going to be worth it? Let me know your thoughts.

The Zernikes

The plan was to cache these where possible, but also to make sure that they were jax safe so they could be used more easily in the context of other programs. It may have been wasted effort but it seems to work for now. Because lax.while_loop, jits the body_fun I do not recommend re-jit-ing at the top level. vmap is fine, and this was the primary goal as we can model multiple apertures of different sizes in one plane.

I actually think that the while_loop is recompiling the body_fun every time and this is what is making it so slow.

Again please let me know what you think.

Regards

Jordan

[benjaminpope]

Thanks for looking into this.

I think there will not for JWST be any need for nonlinear transformations of the aperture. The thing is rock solid. The small variations we need to calibrate are going to be of order 1%, probably including hexikes on each mirror, one or more micrometeorite dents (localized blobs), a bit of Fresnel to really nail the focus of each plane in the optical train, and linear rotations/shears/scalings of parts of the optical train.

I wouldn't spend much more time really nailing jitted recalculation of basis vectors. If they are nearly orthonormal that is almost as good as if they are properly orthonormal - really, you can use any modal basis you like, so long as it spans the aberrations that occur.

[Jordan-Dennis]

Ok cool.
I've been polishing it up this morning and I'm almost certain that there is not way to safely jit the Zernikes except with the while_loop which is slower anyway.

Ironically the problem is this simple mathematical expression:

because the summation limit is dependent on n and m. There are other ways to calculate the radial Zernikes but they are all recursive and jax has no support for recursion.

I've implemented a newer version that does not jit, but which is jax safe for vmap and runs in comparable time to the numpy python implementation. The internal vectorization is pretty neat already so it should run very well on GPU.

Since, we are not planning to do non-linear transformations this is unnecessary for the most part. I will finish with the linear transformations and then update everyone.

[benjaminpope]

Yep, I don't think there is a good way to jit this and do it at runtime. It is better to cache whichever set of modes we choose at __init__ and only optimize over the coefficients.

The higher priority is getting it running for a reasonably idealized JWST, making sure the hexikes exactly match the segments.