New inverse stream map to accelerate convergence

[unalmis]

Inverse stream maps

• Significantly improves convergence of inverse stream maps. Resolves Choose stream s.t. the map to the moving mesh of field lines is trivial #1922 to my satisfaction.
• Truncation rule to remove aliasing error at the shortest wavelengths where the signal to noise ratio is lowest.
• Resolves Increase Y parameter for Heliotron (NFP=19) in effective ripple tutorial #1928 and increases default objective resolution for Y_B.
• Method to plot frequency spectrum of inverse stream map in field line coordinates.

Other performance improvements

• Check-pointing to increase speed and reduce memory consumption of reverse mode differentiation Checkpointing to reduce reverse mode AD memory usage #1347.
• Adds low_ram mode which is same speed and less memory for objective.compute, but slower for objective.grad since JAX is poor at iterative algorithms.
• Fully resolves Memory regression in bounce integrals #1864 by avoiding materialization of a large tensor in memory. Previously, we had closed the issue by adding nuffts as a workaround. The improvement here actually solves the JAX regression.
• Reuses some computations in identifying bounce points to make more efficient.
• Various small improvements to increase cache hits, fusing, and reduce floating point error.
• Transforms an improper field line integral to one on a compact domain where the integrand is periodic to achieve faster convergence.
• Improves performance complexity of interp_to_argmin for Bounce2D from h^4 to spectral.
• Resolves Use OOP for surface integrals with faster methods for tensor product grids #1389 (the non-cosmetic part)

Usability

• Updates plotting utilities to accept kwargs as was needed for my article.
• Method to compute integrals over plasma volume and velocity space in batches is now added to the public API. People kept asking how to use this private function and also used it for their work, so better to make it public.
• Method to make adding new bounce integral objectives simpler. Private for developers only.
• Unifies default parameter setup to resolve Add comments about recommeded resolutions to docstrings #1898 for bounce integrals.
• Resolves Move some general interpolation functions into interpax #1388.

Bugs

• Fixes bug in test_compute_everything.
• Fixes inverse stream map convergence.
• spectral_reverse_bounce_supplement.pdf

Benchmarks

Just go to #2026 and run effective_ripple_profile.py. You will see the large performance improvement from master. The CI benchmarks do not reveal this because those benchmarks are essentially just noise. Note that, using the same parameter inputs, the resolution of this branch is also higher than master due to the faster convergence.

• If you set use_bounce1d=True on that script, you will run out of memory as expected since it is an inferior approach (as expected, you get the OOM in the jacobian before you compute a single bounce integral).
• If you set nufft_eps=0, you need 175 GB to run that script on master (you'll get an OOM and JAX will tell you it needs 175GB), but only 35 GB on this branch.
• Using nuffts, the script requires only 6.7 GB on this branch.

Examples

HELIOTRON

Master branch

This branch

W7-X

Master branch

This branch

NCSX

Master branch

This branch

Removal of spectral aliasing

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    5.94 %    |     3.826e+03      |     4.053e+03      |    227.34    |       39.46        |       36.44        |
  test_proximal_jac_w7x_with_eq_update   |   -0.22 %    |     6.493e+03      |     6.479e+03      |    -14.16    |       162.81       |       160.60       |
  test_proximal_freeb_jac                |   -0.09 %    |     1.321e+04      |     1.320e+04      |    -11.76    |       84.51        |       82.10        |
  test_proximal_freeb_jac_blocked        |   -0.67 %    |     7.528e+03      |     7.478e+03      |    -50.43    |       73.78        |       73.06        |
  test_proximal_freeb_jac_batched        |    0.52 %    |     7.451e+03      |     7.490e+03      |    38.51     |       73.08        |       72.86        |
  test_proximal_jac_ripple               |   -4.61 %    |     3.531e+03      |     3.368e+03      |   -162.74    |       62.58        |       65.67        |
  test_proximal_jac_ripple_bounce1d      |   -0.35 %    |     3.583e+03      |     3.570e+03      |    -12.63    |       74.62        |       76.51        |
  test_eq_solve                          |    1.29 %    |     1.992e+03      |     2.017e+03      |    25.64     |       94.31        |       94.11        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

addressed request

[unalmis]

I assume there is no math error...

Yes, and I have resolved your other comments concerning cosmetics in the code comments.