fourier bounce

[unalmis]

Resolves #1045. Uses pseudospectral technique.

Summary

This solves the dominant cost of optimization objectives relying on the methods used in #854: interpolating DESC's 3D transforms to an optimization-step dependent grid that is dense enough for function approximation with local splines. This is in some sense referred to as off-grid interpolation in literature; it is often a bottleneck.

• The function approximation done here requires DESC transforms on a fixed grid with typical resolution, using FFTs to compute the map $\alpha, \zeta \mapsto \theta(\alpha, \zeta)$ between coordinate systems. This enables evaluating functions along field lines without root-finding.
• The faster convergence of spectral interpolation requires a less dense grid to interpolate onto from DESC's 3D transforms.
• Spectral function approximation is more accurate than cubic splines.
• 2D interpolation enables tracing the field line for more toroidal transits.
• With these improvements, optimization with objectives on ε and Γ_c with Fourier Bounce #1290 now take 6gb of memory whereas before they may not have fit onto today's gpu.

Future performance improvements

• Improve coordinate mapping performance #1154
• Nonuniform fast transforms #1294
• Patch for differentiable code with dynamic shapes #1303
• Upsample data above midplane to full grid assuming stellarator symmetry #1206
• Optimizers/objectives with auxiliary output #1034

[codecov]

Codecov Report

Attention: Patch coverage is 96.60194% with 21 lines in your changes missing coverage. Please review.

Project coverage is 95.53%. Comparing base (005c4f8) to head (9250166).
Report is 121 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/integrals/basis.py	91.77%	13 Missing ⚠️
desc/integrals/bounce_integral.py	95.62%	7 Missing ⚠️
desc/integrals/_bounce_utils.py	98.83%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1119      +/-   ##
==========================================
- Coverage   95.59%   95.53%   -0.07%     
==========================================
  Files          96       96              
  Lines       24601    25044     +443     
==========================================
+ Hits        23518    23926     +408     
- Misses       1083     1118      +35     

Files with missing lines	Coverage Δ
desc/backend.py	90.44% <100.00%> (ø)
desc/compute/_basis_vectors.py	100.00% <100.00%> (ø)
desc/compute/_core.py	100.00% <100.00%> (ø)
desc/compute/_metric.py	100.00% <100.00%> (ø)
desc/equilibrium/equilibrium.py	96.07% <ø> (ø)
desc/grid.py	94.58% <100.00%> (+0.02%)	⬆️
desc/integrals/__init__.py	100.00% <100.00%> (ø)
desc/integrals/_interp_utils.py	100.00% <100.00%> (ø)
desc/integrals/_quad_utils.py	100.00% <100.00%> (ø)
desc/transform.py	93.29% <100.00%> (ø)
... and 4 more

... and 2 files with indirect coverage changes

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     +0.65 +/- 4.41     | +3.84e-03 +/- 2.62e-02 |  5.99e-01 +/- 2.5e-02  |  5.95e-01 +/- 6.6e-03  |
 test_build_transform_fft_highres        |     +0.25 +/- 1.18     | +2.39e-03 +/- 1.13e-02 |  9.56e-01 +/- 7.8e-03  |  9.53e-01 +/- 8.1e-03  |
 test_equilibrium_init_lowres            |     +0.20 +/- 0.87     | +7.41e-03 +/- 3.26e-02 |  3.78e+00 +/- 2.2e-02  |  3.77e+00 +/- 2.4e-02  |
 test_objective_compile_atf              |     +0.23 +/- 4.25     | +1.84e-02 +/- 3.33e-01 |  7.85e+00 +/- 2.5e-01  |  7.84e+00 +/- 2.2e-01  |
 test_objective_compute_atf              |     -1.35 +/- 1.99     | -1.43e-04 +/- 2.10e-04 |  1.04e-02 +/- 1.9e-04  |  1.06e-02 +/- 9.2e-05  |
 test_objective_jac_atf                  |     +2.72 +/- 4.06     | +5.05e-02 +/- 7.55e-02 |  1.91e+00 +/- 4.8e-02  |  1.86e+00 +/- 5.9e-02  |
 test_perturb_1                          |     +0.84 +/- 2.68     | +1.17e-01 +/- 3.71e-01 |  1.40e+01 +/- 6.7e-02  |  1.38e+01 +/- 3.7e-01  |
 test_proximal_jac_atf                   |     -0.09 +/- 1.73     | -7.58e-03 +/- 1.41e-01 |  8.16e+00 +/- 1.2e-01  |  8.17e+00 +/- 7.1e-02  |
 test_proximal_freeb_compute             |     +0.00 +/- 0.76     | +4.87e-06 +/- 1.50e-03 |  1.96e-01 +/- 9.9e-04  |  1.96e-01 +/- 1.1e-03  |
 test_solve_fixed_iter_compiled          |     +0.42 +/- 1.29     | +7.01e-02 +/- 2.16e-01 |  1.68e+01 +/- 1.2e-01  |  1.67e+01 +/- 1.8e-01  |
 test_build_transform_fft_lowres         |     +4.63 +/- 5.28     | +2.46e-02 +/- 2.80e-02 |  5.56e-01 +/- 2.2e-02  |  5.31e-01 +/- 1.7e-02  |
 test_equilibrium_init_medres            |     +2.03 +/- 4.20     | +8.68e-02 +/- 1.79e-01 |  4.36e+00 +/- 3.9e-02  |  4.27e+00 +/- 1.8e-01  |
 test_equilibrium_init_highres           |     +2.74 +/- 1.96     | +1.50e-01 +/- 1.07e-01 |  5.63e+00 +/- 5.9e-02  |  5.48e+00 +/- 8.9e-02  |
 test_objective_compile_dshape_current   |     +1.03 +/- 1.34     | +3.99e-02 +/- 5.22e-02 |  3.92e+00 +/- 2.5e-02  |  3.88e+00 +/- 4.6e-02  |
 test_objective_compute_dshape_current   |     -0.16 +/- 1.46     | -5.98e-06 +/- 5.36e-05 |  3.67e-03 +/- 4.6e-05  |  3.67e-03 +/- 2.7e-05  |
 test_objective_jac_dshape_current       |     +3.11 +/- 8.49     | +1.25e-03 +/- 3.43e-03 |  4.16e-02 +/- 2.5e-03  |  4.04e-02 +/- 2.4e-03  |
 test_perturb_2                          |     -0.26 +/- 3.55     | -5.09e-02 +/- 6.93e-01 |  1.94e+01 +/- 6.4e-01  |  1.95e+01 +/- 2.6e-01  |
 test_proximal_freeb_jac                 |     +0.20 +/- 1.39     | +1.51e-02 +/- 1.04e-01 |  7.49e+00 +/- 8.3e-02  |  7.47e+00 +/- 6.3e-02  |
 test_solve_fixed_iter                   |     +1.50 +/- 2.19     | +4.29e-01 +/- 6.26e-01 |  2.90e+01 +/- 5.6e-01  |  2.85e+01 +/- 2.8e-01  |
 test_LinearConstraintProjection_build   |     -0.04 +/- 2.31     | -9.80e-03 +/- 5.26e-01 |  2.27e+01 +/- 4.8e-01  |  2.27e+01 +/- 2.2e-01  |

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[unalmis]

Re requesting review since the requested changes were made, but you are welcome to delay as this is not urgent

[rahulgaur104]

Merge if all tests pass.

[f0uriest]

Also it looks like #1290 makes some changes/updates to the Bounce2D apis, would be good to do all of that here

[unalmis]

Also it looks like #1290 makes some changes/updates to the Bounce2D apis, would be good to do all of that here

Done. The only user facing change is to accept a dictionary instead of a list in bounce.integrate so that the compute function is

def _thing(data, B, pitch):
       return B * pitch * data["thing"] * data["things2"]

[rahulgaur104]

Tests not passing right now. Merge once all tests pass.