Nonuniform fast transforms

[unalmis]

In general, the computation of locally defined or singular quantities will be more efficient with non-uniform fast transforms. In particular, I am referring to the type 2 nonuniform fast transform, which evaluates spectral coefficients at non-uniformly spaced points.

Although a fixed grid with nodes uniformly spaced throughout the plasma is fine for optimizing continuous functions defined throughout the volume, it is of course not well-suited for locally defined quantities such as bounce integrals. There are two hurdles with computing these quantities efficiently:

1. Require computation along field lines.
   • fourier bounce #1119 should cover this, improving upon Bounce averaging #854.
2. They are defined by singular integrals.
   • This requires evaluating the function at points concentrated near the singularity. For flux surface average of bounce integrals in particular, there are infinitely many integrals to be done, so regardless of the efficiency of the quadrature, there will be many function evaluations. Hence, linearithmic performance of fast non-uniform transform is desired over the quadratic performance of MMTs.

We should pick a JAX compatible NFFT package and write drop in replacements for the currently used MMTs. We can check the $\epsilon$ errors of these external NFFT libraries by comparing to those utilities on the tests in that PR. (Note that for strongly singular integrals $1 / \sqrt(1 - \lambda |B|)$, we will maybe still use MMT's for |B| since the integrands are not smooth functions of |B|, and the nodes are packed more tightly near the edges. This should be fine since |B| is parameterized by a series expansion in a single variable along each cut of the field line, so it is cheaper. The other quantities in the integrand are 2D functions where this would be more useful. For the weakly singular ones, all the quantities should be able to use NFFTs without issue).

Here are some candidates:

1. https://github.com/jakevdp/nfft
2. https://github.com/flatironinstitute/jax-finufft

[acerfon]

For what it's worth, we have been very happy with the Flatiron's NUFFT, which we have used with much success for a new Particle in Fourier scheme:
https://www.sciencedirect.com/science/article/pii/S0021999124006387#:~:text=Journal%20of%20Computational%20Physics.%20Volume

I highly recommend that https://github.com/flatironinstitute/jax-finufft library!

[acerfon]

As another 2-cents, for the singular integrals you are talking about, I very highly recommend the Kapur-Rokhlin quadrature schemes, which we have also found easy to use, and super accurate, in our case for the evaluation of the virtual casing principle:
https://www.jstor.org/stable/2952054

[f0uriest]

I've also used finufft and really like it, I thought about using it for some of the stuff in #725, the main reason I didn't was that they don't have prebuilt cuda wheels so installing it with gpu support can be annoying. It looks like they have an open issue for that (flatironinstitute/jax-finufft#88) so if that can be fixed it would be great for us.

NUFFTs may also make for #641 easier for some cases.