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.

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.
• Increase cache hits, fusing, and reduce floating point error in computing bounce points (very important for accurate integrals).
• 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 fourth order to spectral as required for Alpert quadrature.
• Resolves Use OOP for surface integrals with faster methods for tensor product grids #1389.

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.

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 34 GB on this branch.
• Using nuffts, the script requires only 6.5 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

Dynamic shapes

• resolves Patch for differentiable code with dynamic shapes #1303 .
• Activates newton step to find bounce points.
  • It can be shown there exists O(sqrt(epsilon)) error for bounce integrals with 1/v_ll where epsilon is error of bounce point. For v_|| integrals error is conveniently O(epsilon times sqrt(epsilon)). Hence, the spline method would require thousands of knots per transit for just a couple digits of accuracy, and it would stop convergence at epsilon<=1e-5 (so sqrt(epsilon) <=3 digits) error due to condition number. Of course, the spline=False method has always computed the points with spectral accuracy and has very fast convergence after New inverse stream map to accelerate convergence #1919 ; that method converges to epsilon = machine precision without the newton step. With the newton step, fast convergence is achieved with the spline method as well.
  • I suspect fast ion confinement optimization will be easier now.

Benchmarks

Here is a timing benchmark on my CPU with nufft_eps=1e-6. Prior to this PR, every adjoint call to nufft1 took >= 1 second and the full computation was 34 seconds. Now every adjoint call to nufft1 is 250 milliseconds, and the full computation is 14 seconds. These improvements become larger as the sparsity grows and error tolerance parameter for the nuffts epsilon tends to 0. Likewise, the improvement grows linearly with the problem size. As this is called within an optimization loop where time and memory are tight, the improvement is significant.

Before

After

[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                 |    1.38 %    |     4.202e+03      |     4.260e+03      |    58.14     |       41.77        |       38.76        |
  test_proximal_jac_w7x_with_eq_update   |   -2.29 %    |     6.584e+03      |     6.433e+03      |   -151.00    |       158.32       |       157.41       |
  test_proximal_freeb_jac                |   -0.19 %    |     1.340e+04      |     1.338e+04      |    -24.80    |       86.52        |       87.00        |
  test_proximal_freeb_jac_blocked        |    0.53 %    |     7.736e+03      |     7.777e+03      |    41.04     |       76.38        |       76.48        |
  test_proximal_freeb_jac_batched        |    0.50 %    |     7.669e+03      |     7.707e+03      |    38.25     |       75.87        |       77.10        |
  test_proximal_jac_ripple               |   -3.71 %    |     3.733e+03      |     3.594e+03      |   -138.53    |       61.75        |       65.36        |
  test_proximal_jac_ripple_bounce1d      |    1.89 %    |     3.855e+03      |     3.927e+03      |    72.68     |       74.60        |       76.63        |
  test_eq_solve                          |    0.23 %    |     2.222e+03      |     2.227e+03      |     5.07     |       98.03        |       97.03        |

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

addressed request

[unalmis]

Someone else who is still working on desc will need to maintain this PR from now on.

[unalmis]

@dpanici @YigitElma @f0uriest @ddudt @rahulgaur104 . I requested an ETA 3 months ago for when this would be merged, but no one has replied.

[YigitElma]

I don't like having a new dependency. We already have many problems with JAX and related packages. I don't feel comfortable approving the current state of the PR.

[unalmis]

• The dependency uses the same packages that desc already relies on. It's not going to cause problems. Your issue is with JAX, not my code.
• I had predicted the recent jax finufft issue and put guards in place to avoid them in future. When those guards were explicitly removed by other developers without my approval, a consequence was that issue. That would have been avoided if my code had not been modified.
• The other developers requested I move it to an external package. I waited 3 months to make the change (October to December) just in case someone changed their mind. I even
  confirmed it again via a poll one week prior to when I added the dependency. To now state that you are blocking the work because of that dependency is unreasonable , especially given the thorough lengths I went to to avoid this particular issue.

[YigitElma]

• To now state that you are blocking the work because of that dependency is unreasonable

I am not blocking it. If others approve it, you can merge it; that is fine. I don't have to approve all PRs. I've reviewed it multiple times and spent a couple of days on it.

[dpanici]

@dpanici @YigitElma @f0uriest @ddudt @rahulgaur104 . I requested an ETA 3 months ago for when this would be merged, but no one has replied.

Sorry for the delay on my part, I will review this in the next week. Hopefully by the end of this week once my thesis is turned in. I need to also see what is exactly going on with the finufft package, maybe we can set an upper bound to avoid the breaking changes they introduce in the latest release