fourier bounce

[unalmis]

Resolves #1045. Uses pseudospectral technique with smarter mathematics for bounce integration. Needed for neoclassical optimization.

High level summary

The improvements are discussed more in the code documentation. This solves the dominant cost of optimization objectives relying on the methods used in #854: interpolating DESC transforms to an optimization-step dependent grid that was dense enough for local splines to reconstruct the field line. This is sometimes referred to as off-grid interpolation in literature; it is often a bottleneck. (Note that the bounce integrals (number of pitch angles etc.) didn't even show up in memory profiling for Bounce1D).

The function approximation done here requires DESC transforms on a fixed grid with typical resolution, using FFTs to parameterize the spectral series expansion of quantities in field line coordinates. Evaluation of the FFT's is made less expensive by parameterizing the 2d function by one variable along field lines (sort of like partial summation). (At a basic level, the faster convergence of spectral interpolation requires a less dense grid to interpolate onto, and the 2D function approximation allows for larger number of toroidal transits to be integrated over.).

The next steps are

1. using non-uniform ffts for interpolation to quadrature points. (this should be 1 line change once we pick a good library. see the github issue for candidates)
2. improving the root finding algorithm (instructions to implement are in the code). (The Bounce1D root finding is fine, when we move to pseudospectral though it's best to change the algorithm).

Tasks

• write algorithm
• document algorithm
• cross-reference results of effective ripple and Gamma_c with Bounce1D method as tests; forwarded task to Writing ε and Γ_c compute funs with new methods #1290

[codecov]

Codecov Report

Attention: Patch coverage is 92.11957% with 29 lines in your changes missing coverage. Please review.

Project coverage is 95.41%. Comparing base (e35436d) to head (60fc00f).

Files with missing lines	Patch %	Lines
desc/integrals/basis.py	90.85%	15 Missing ⚠️
desc/integrals/bounce_integral.py	87.38%	14 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1119      +/-   ##
==========================================
- Coverage   95.45%   95.41%   -0.04%     
==========================================
  Files          96       96              
  Lines       23806    24160     +354     
==========================================
+ Hits        22723    23052     +329     
- Misses       1083     1108      +25     

Files with missing lines	Coverage Δ
desc/grid.py	94.53% <ø> (ø)
desc/integrals/__init__.py	100.00% <100.00%> (ø)
desc/integrals/bounce_utils.py	93.33% <100.00%> (+0.13%)	⬆️
desc/integrals/interp_utils.py	100.00% <100.00%> (ø)
desc/integrals/quad_utils.py	100.00% <ø> (ø)
desc/transform.py	93.27% <100.00%> (ø)
desc/utils.py	92.19% <100.00%> (+1.36%)	⬆️
desc/integrals/bounce_integral.py	91.25% <87.38%> (-8.75%)	⬇️
desc/integrals/basis.py	91.70% <90.85%> (-5.07%)	⬇️

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |    +12.27 +/- 15.44    | +7.77e-02 +/- 9.79e-02 |  7.11e-01 +/- 9.7e-02  |  6.34e-01 +/- 1.5e-02  |
 test_build_transform_fft_highres        |     -0.82 +/- 5.86     | -8.73e-03 +/- 6.25e-02 |  1.06e+00 +/- 4.6e-02  |  1.07e+00 +/- 4.3e-02  |
 test_equilibrium_init_lowres            |     +7.57 +/- 7.46     | +3.14e-01 +/- 3.09e-01 |  4.46e+00 +/- 2.6e-01  |  4.15e+00 +/- 1.6e-01  |
 test_objective_compile_atf              |     +2.36 +/- 12.18    | +1.93e-01 +/- 9.98e-01 |  8.39e+00 +/- 7.2e-01  |  8.20e+00 +/- 7.0e-01  |
 test_objective_compute_atf              |     -1.57 +/- 3.04     | -1.65e-04 +/- 3.18e-04 |  1.03e-02 +/- 1.8e-04  |  1.05e-02 +/- 2.6e-04  |
 test_objective_jac_atf                  |     +2.26 +/- 3.02     | +4.43e-02 +/- 5.93e-02 |  2.01e+00 +/- 2.1e-02  |  1.96e+00 +/- 5.5e-02  |
 test_perturb_1                          |     +2.04 +/- 3.22     | +2.59e-01 +/- 4.09e-01 |  1.30e+01 +/- 2.3e-01  |  1.27e+01 +/- 3.4e-01  |
 test_proximal_jac_atf                   |     +0.57 +/- 1.62     | +4.72e-02 +/- 1.33e-01 |  8.28e+00 +/- 1.2e-01  |  8.24e+00 +/- 5.2e-02  |
 test_proximal_freeb_compute             |     +1.02 +/- 2.20     | +1.88e-03 +/- 4.07e-03 |  1.87e-01 +/- 3.6e-03  |  1.85e-01 +/- 1.8e-03  |
 test_build_transform_fft_lowres         |     -0.56 +/- 4.42     | -3.03e-03 +/- 2.39e-02 |  5.38e-01 +/- 1.8e-02  |  5.41e-01 +/- 1.6e-02  |
 test_equilibrium_init_medres            |     +0.22 +/- 2.25     | +9.33e-03 +/- 9.57e-02 |  4.26e+00 +/- 9.0e-02  |  4.25e+00 +/- 3.4e-02  |
 test_equilibrium_init_highres           |     +2.07 +/- 2.00     | +1.16e-01 +/- 1.12e-01 |  5.73e+00 +/- 9.7e-02  |  5.61e+00 +/- 5.7e-02  |
 test_objective_compile_dshape_current   |     -0.08 +/- 1.78     | -2.98e-03 +/- 6.97e-02 |  3.91e+00 +/- 4.0e-02  |  3.91e+00 +/- 5.7e-02  |
 test_objective_compute_dshape_current   |     +3.90 +/- 2.50     | +1.35e-04 +/- 8.64e-05 |  3.59e-03 +/- 5.7e-05  |  3.46e-03 +/- 6.5e-05  |
 test_objective_jac_dshape_current       |     +2.03 +/- 7.57     | +8.23e-04 +/- 3.08e-03 |  4.15e-02 +/- 2.5e-03  |  4.06e-02 +/- 1.8e-03  |
 test_perturb_2                          |     -0.11 +/- 1.42     | -2.05e-02 +/- 2.57e-01 |  1.80e+01 +/- 1.4e-01  |  1.81e+01 +/- 2.1e-01  |
 test_proximal_freeb_jac                 |     +0.10 +/- 0.96     | +7.70e-03 +/- 7.08e-02 |  7.34e+00 +/- 5.2e-02  |  7.34e+00 +/- 4.8e-02  |
 test_solve_fixed_iter                   |     +0.44 +/- 60.81    | +2.20e-02 +/- 3.06e+00 |  5.05e+00 +/- 2.1e+00  |  5.03e+00 +/- 2.2e+00  |

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[unalmis]

This seems intuitive to me, so I didn't formally prove it. A "proof by example" would be to plot the Fourier Chebyshev series $\theta(\alpha(\vartheta, \zeta), \zeta)$ along a field line. Now since $\theta$ is continuous and $\alpha$ is continuous in $\vartheta, \zeta$, then by varying $\vartheta, \zeta$ along a field line there is guaranteed to be no discontinuities at the cuts (modulo 2pi). Now check if this matches the double Fourier series $\theta(\vartheta, \zeta)$. You should find that given sufficient convergence, the fourier chebyshev and double fourier series give essentially the same plots, and yet you can always zoom in to the plot $\theta(\alpha, \zeta)$ between different alpha values that specify the same field line to find a discontinuity (it will vanish as the series converges, but you should always be able to see it by zooming in).

[unalmis]

Ok got things working, I think we're good to just write the compute functions for effective ripple and Gamma_c now and see what they give

[unalmis]

commenting for reference

[unalmis]

commenting for reference