FINUFFT `plan` api

[unalmis]

Also an API like FINUFFT where I can call nufft multiple times with the same plan would be convenient to avoid materializing a tensor by concatenating source arrays.

Here is what I mean.

Originally posted by @unalmis in #155 (comment)

[dfm]

Please be more explicit about your specific use case. Presumably you want to use this in eager mode? I'm not totally convinced that this would provide a real performance benefit within a larger JAX program under jit, and it wouldn't be that straightforward to implement, so you'd need to explicitly demonstrate (i.e. with code) that you have a use case that significantly underperforms using the existing APIs which already transparently reuse the plan.

[unalmis]

Thank you for the reply. Here are specific use cases where the plan (which FINUFFT also labels as the guru) interface in FINUFFT is preferred to the existing implementation in jax-finufft.

1. Suppose the non-uniform points x are known a priori and the set Fourier series {c_i | i in (1, ..., n)} are not. Then the plan interface allows computing the nufft for each Fourier series c_i efficiently. In contrast, the existing API requires building the nufft again for each input. This is example where the performance of the plan API is significantly better according to the FINUFFT documentation.
2. Suppose as above, but we are in the more forgiving scenario where the full set {c_i} is known simultaneously. Then the plan interface allows computing each nufft efficiently with data = {c_i : plan.execute(data[c_i], x) for c_i in data}. In contrast, the existing API requires materializing a tensor, then mapping this tensor back to a dictionary. Subsequent computations will then need to index the required subsets out of this tensor. (These expenses are unlikely to be free, as memory is required for the tensor and computation for the indexing). Ignoring that, this is just an example where the plan API is more convenient, leading to simpler user facing code.
3. JAX has bugs. Two relevant ones are GitHub issues #26689 and #27591. Experience suggests that constructing the tensor in 2. would lead to incorrect computations for large tensors. This is an example where the plan API is needed for correctness.
4. In applications with singular integrals, it is common to require evaluation of one of the Fourier series to much higher accuracy than the rest. With plan API I can choose a tighter error eps for that Fourier series compared to the rest. This is again useful for correctness and performance.

[unalmis]

FINUFFT details benchmarks for performance improvements, so I skipped the request for a benchmark.