Use fast transforms in radial direction for Zernike

[unalmis]

Like #641 , but in the radial direction. This should make computing the Zernike polynomials faster, in particular for the simultaneous computation at many points. The main benefit is that it should be cheaper to differentiate through.

[YigitElma]

There was a paper on calculating the radial part of the Zernike polynomials using FFT's but it was a fit, not using the actual definition of Zernike radial Polynomials. @f0uriest had a version of that, but I doubt it would give much benefit.
https://arxiv.org/abs/2304.14492

The poloidal part of Zernike polynomials is just the Fourier series, so it should be easy.

[YigitElma]

Also, I have a general question about map_coordinates and the derivative of Zernike. The way we define the custom derivative of Zernike polynomials is that we only allow it to have a derivative in radial coordinate, $\rho$. But I cannot understand why would we need that derivative for optimization purposes. The only derivative we are taking in optimization is the spectral coefficients, and zernike_radial function doesn't take that as an argument, so the autodiff shouldn't try to differentiate that, as far as I understand. For the eq.solve and optimize that don't use map_coordinates with $\rho, \alpha, \zeta$, I don't think that derivative information is used at all currently.

My point is instead of trying to make it faster, we can try to debug why it is necessary in the first place?

[unalmis]

I marked this low priority, since typically fast transforms are useful.

If we have a function to compute f that evaluates the zernike polynomials on some grid, and the grid depends on the optimization parameters (e.g. spectral coeffcients of the stream functions), and we want to see how f changes wrt to the optimization parameters, then the AD tool should track the derivatives of the zernike computation on that grid in the same way it should track derivatives of the operations that define f.

IIRC there exists a corresponding fast algorithm (DPT) to compute any orthogonal polynomial at a set of points. If this has already been tested it may be worth mentioning somewhere.

[YigitElma]

If we have a function to compute f that evaluates the zernike polynomials on some grid, and the grid depends on the optimization parameters (e.g. spectral coeffcients of the stream functions), and we want to see how f changes wrt to the optimization parameters

Can you give an example of this? I think you refer to optimizations in https://desc-docs.readthedocs.io/en/latest/notebooks/tutorials/EffectiveRipple.html#Optimizing-Heliotron, but it seems like you give the specific $\rho$ values to calculate on. Sorry I had to review it.

[unalmis]

If we have a function to compute f that evaluates the zernike polynomials on some grid, and the grid depends on the optimization parameters (e.g. spectral coeffcients of the stream functions), and we want to see how f changes wrt to the optimization parameters

Can you give an example of this? I think you refer to optimizations in https://desc-docs.readthedocs.io/en/latest/notebooks/tutorials/EffectiveRipple.html#Optimizing-Heliotron, but it seems like you give the specific $\rho$ values to calculate on. Sorry I had to review it.

That quantity requires following field lines to compute; depending on the algorithm implemented to compute such a quantity, the zernike polynomials can be recomputed on an optimization-step dependent grid. Then the computation involves derivatives of fourier zernike series $\sum_{lmn} a_{lmn} R_l(\rho) f_m(\theta) f_n(\zeta)$. We have yet to implement partial summation to avoid recomputing the radial term, so the radial term always gets recomputed even if $\rho$ never changes. I'm sure there also exists quantities of interest where the radial coordinate is not fixed, perhaps tracking radial drift for ep's etc.

It's basically just 641 but in radial direction so similar benefit

[unalmis]

feel free to edit the description of the issue if there's something to fix

[YigitElma]

Also related to #1243