ExternalObjective function to wrap external codes

[ddudt]

Creates an abstract base class for wrapping external codes with finite differences, like GX, TERPSICHORE, etc.

TODO:

• vectorize
• reverse mode gradient
• tests
• tutorial

[codecov]

Codecov Report

Attention: Patch coverage is 96.10390% with 3 lines in your changes missing coverage. Please review.

Project coverage is 95.61%. Comparing base (34dbac0) to head (96aec58).

Files with missing lines	Patch %	Lines
desc/backend.py	75.00%	2 Missing ⚠️
desc/objectives/_generic.py	97.50%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1028   +/-   ##
=======================================
  Coverage   95.60%   95.61%           
=======================================
  Files          98       98           
  Lines       25417    25485   +68     
=======================================
+ Hits        24300    24367   +67     
- Misses       1117     1118    +1     

Files with missing lines	Coverage Δ
desc/integrals/surface_integral.py	100.00% <ø> (ø)
desc/objectives/__init__.py	100.00% <100.00%> (ø)
desc/optimize/least_squares.py	99.33% <100.00%> (ø)
desc/utils.py	93.88% <100.00%> (+0.50%)	⬆️
desc/objectives/_generic.py	97.54% <97.50%> (-0.06%)	⬇️
desc/backend.py	90.16% <75.00%> (-0.29%)	⬇️

... and 2 files with indirect coverage changes

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |    +10.00 +/- 4.18     | +5.35e-02 +/- 2.24e-02 |  5.88e-01 +/- 2.0e-02  |  5.35e-01 +/- 9.9e-03  |
-test_equilibrium_init_medres            |    +10.68 +/- 1.71     | +4.44e-01 +/- 7.11e-02 |  4.60e+00 +/- 6.6e-02  |  4.16e+00 +/- 2.6e-02  |
 test_equilibrium_init_highres           |     +3.54 +/- 4.11     | +1.94e-01 +/- 2.25e-01 |  5.66e+00 +/- 2.2e-01  |  5.47e+00 +/- 4.2e-02  |
 test_objective_compile_dshape_current   |     +0.59 +/- 1.08     | +2.31e-02 +/- 4.25e-02 |  3.95e+00 +/- 3.4e-02  |  3.93e+00 +/- 2.6e-02  |
 test_objective_compute_dshape_current   |     -0.97 +/- 1.85     | -5.01e-05 +/- 9.53e-05 |  5.10e-03 +/- 4.0e-05  |  5.15e-03 +/- 8.7e-05  |
 test_objective_jac_dshape_current       |     -3.77 +/- 6.73     | -1.65e-03 +/- 2.95e-03 |  4.22e-02 +/- 2.1e-03  |  4.38e-02 +/- 2.1e-03  |
 test_perturb_2                          |     +1.63 +/- 4.20     | +3.23e-01 +/- 8.29e-01 |  2.01e+01 +/- 8.1e-01  |  1.98e+01 +/- 2.0e-01  |
 test_proximal_freeb_jac                 |     -0.27 +/- 1.12     | -1.98e-02 +/- 8.24e-02 |  7.37e+00 +/- 7.1e-02  |  7.39e+00 +/- 4.2e-02  |
 test_solve_fixed_iter                   |     +0.79 +/- 2.29     | +2.54e-01 +/- 7.32e-01 |  3.22e+01 +/- 4.0e-01  |  3.20e+01 +/- 6.1e-01  |
 test_LinearConstraintProjection_build   |     -0.15 +/- 2.86     | -1.57e-02 +/- 2.96e-01 |  1.03e+01 +/- 2.1e-01  |  1.04e+01 +/- 2.1e-01  |
 test_build_transform_fft_midres         |     -0.26 +/- 5.51     | -1.58e-03 +/- 3.40e-02 |  6.16e-01 +/- 1.9e-02  |  6.18e-01 +/- 2.8e-02  |
 test_build_transform_fft_highres        |     +0.96 +/- 2.66     | +9.40e-03 +/- 2.60e-02 |  9.86e-01 +/- 1.8e-02  |  9.76e-01 +/- 1.9e-02  |
 test_equilibrium_init_lowres            |     +0.34 +/- 3.07     | +1.33e-02 +/- 1.20e-01 |  3.91e+00 +/- 1.0e-01  |  3.90e+00 +/- 6.1e-02  |
 test_objective_compile_atf              |     +0.78 +/- 0.86     | +6.42e-02 +/- 7.05e-02 |  8.24e+00 +/- 4.4e-02  |  8.18e+00 +/- 5.5e-02  |
 test_objective_compute_atf              |     -3.24 +/- 6.66     | -5.34e-04 +/- 1.10e-03 |  1.59e-02 +/- 3.4e-04  |  1.65e-02 +/- 1.0e-03  |
 test_objective_jac_atf                  |     +2.22 +/- 1.75     | +4.31e-02 +/- 3.41e-02 |  1.99e+00 +/- 2.6e-02  |  1.94e+00 +/- 2.2e-02  |
 test_perturb_1                          |     +0.56 +/- 2.32     | +8.44e-02 +/- 3.48e-01 |  1.50e+01 +/- 3.1e-01  |  1.50e+01 +/- 1.6e-01  |
 test_proximal_jac_atf                   |     -0.35 +/- 1.26     | -2.88e-02 +/- 1.05e-01 |  8.25e+00 +/- 5.8e-02  |  8.28e+00 +/- 8.7e-02  |
 test_proximal_freeb_compute             |     -0.23 +/- 1.40     | -4.72e-04 +/- 2.83e-03 |  2.01e-01 +/- 1.3e-03  |  2.02e-01 +/- 2.5e-03  |
 test_solve_fixed_iter_compiled          |     -0.20 +/- 1.57     | -4.26e-02 +/- 3.27e-01 |  2.07e+01 +/- 8.2e-02  |  2.08e+01 +/- 3.2e-01  |

[ddudt]

This takes a minute or so to run on my laptop, since VMECIO.save is slow. If we really need to speed this up we could either reduce the equilibrium resolution or manually save only the VMEC quantities that are used.

[ddudt]

I sped this up significantly by manually saving the few quantities instead of calling VMECIO.save

[f0uriest]

Might be useful when we want to do multithreading jax-ml/jax#24756

[ddudt]

I'm having an issue with the new test in this PR after updating with master where the code gets stuck in some infinite loop. I can make the test pass locally if I add a print statement after the optimizer's first Jacobian call -- that seems to snap it back somehow and then it continues to run fine. Any ideas of what could be going on? This isn't using any multiprocessing

[dpanici]

Put wrapper fxn in backend which checks JAX version and uses correct API

[dpanici]

may want io_callback instead

[ddudt]

This works with JAX v0.4.35 and later. Need to add backwards compatibility.

I think any other vmap_method besides "expand_dims" would also work since it only changes the dimensions for internal calculations and not the final result.