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.05263% with 3 lines in your changes missing coverage. Please review.

Project coverage is 95.32%. Comparing base (225d44e) to head (37f9ee3).

Files	Patch %	Lines
desc/backend.py	75.00%	2 Missing ⚠️
desc/objectives/_generic.py	97.43%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1028      +/-   ##
==========================================
- Coverage   95.33%   95.32%   -0.01%     
==========================================
  Files          90       90              
  Lines       22643    22711      +68     
==========================================
+ Hits        21586    21650      +64     
- Misses       1057     1061       +4     

Files	Coverage Δ
desc/integrals/surface_integral.py	100.00% <ø> (ø)
desc/objectives/__init__.py	100.00% <100.00%> (ø)
desc/utils.py	92.88% <100.00%> (+0.72%)	⬆️
desc/objectives/_generic.py	97.46% <97.43%> (-0.06%)	⬇️
desc/backend.py	89.94% <75.00%> (-0.31%)	⬇️

... and 3 files with indirect coverage changes

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     -2.57 +/- 4.98     | -1.64e-02 +/- 3.19e-02 |  6.24e-01 +/- 1.9e-02  |  6.40e-01 +/- 2.5e-02  |
 test_build_transform_fft_highres        |     +2.56 +/- 10.35    | +2.62e-02 +/- 1.06e-01 |  1.05e+00 +/- 1.0e-01  |  1.02e+00 +/- 1.5e-02  |
 test_equilibrium_init_lowres            |     -1.41 +/- 4.68     | -5.69e-02 +/- 1.89e-01 |  3.99e+00 +/- 1.1e-01  |  4.05e+00 +/- 1.5e-01  |
 test_objective_compile_atf              |     +0.10 +/- 3.74     | +7.97e-03 +/- 2.98e-01 |  7.98e+00 +/- 2.3e-01  |  7.97e+00 +/- 1.9e-01  |
 test_objective_compute_atf              |     -4.70 +/- 4.00     | -5.24e-04 +/- 4.47e-04 |  1.06e-02 +/- 9.2e-05  |  1.12e-02 +/- 4.4e-04  |
 test_objective_jac_atf                  |     -2.85 +/- 2.12     | -5.68e-02 +/- 4.24e-02 |  1.94e+00 +/- 2.7e-02  |  2.00e+00 +/- 3.3e-02  |
 test_perturb_1                          |     +0.12 +/- 2.38     | +1.53e-02 +/- 3.14e-01 |  1.32e+01 +/- 2.6e-01  |  1.32e+01 +/- 1.7e-01  |
 test_proximal_jac_atf                   |     +0.25 +/- 0.86     | +2.02e-02 +/- 7.03e-02 |  8.18e+00 +/- 3.9e-02  |  8.16e+00 +/- 5.8e-02  |
 test_proximal_freeb_compute             |     -0.07 +/- 0.85     | -1.36e-04 +/- 1.56e-03 |  1.83e-01 +/- 1.2e-03  |  1.83e-01 +/- 1.0e-03  |
 test_build_transform_fft_lowres         |     -1.67 +/- 10.30    | -8.77e-03 +/- 5.40e-02 |  5.16e-01 +/- 4.9e-02  |  5.25e-01 +/- 2.3e-02  |
 test_equilibrium_init_medres            |     +0.96 +/- 1.10     | +3.84e-02 +/- 4.42e-02 |  4.05e+00 +/- 4.0e-02  |  4.02e+00 +/- 1.9e-02  |
 test_equilibrium_init_highres           |     +0.79 +/- 4.58     | +4.28e-02 +/- 2.47e-01 |  5.44e+00 +/- 2.1e-01  |  5.39e+00 +/- 1.3e-01  |
 test_objective_compile_dshape_current   |     +0.18 +/- 1.08     | +6.94e-03 +/- 4.10e-02 |  3.79e+00 +/- 1.3e-02  |  3.79e+00 +/- 3.9e-02  |
 test_objective_compute_dshape_current   |     +0.97 +/- 1.91     | +3.53e-05 +/- 6.92e-05 |  3.66e-03 +/- 4.4e-05  |  3.62e-03 +/- 5.3e-05  |
 test_objective_jac_dshape_current       |     +1.00 +/- 5.27     | +3.98e-04 +/- 2.10e-03 |  4.02e-02 +/- 1.7e-03  |  3.98e-02 +/- 1.2e-03  |
 test_perturb_2                          |     +0.23 +/- 0.93     | +3.92e-02 +/- 1.61e-01 |  1.73e+01 +/- 1.2e-01  |  1.73e+01 +/- 1.1e-01  |
 test_proximal_freeb_jac                 |     +0.59 +/- 0.97     | +4.32e-02 +/- 7.10e-02 |  7.37e+00 +/- 6.3e-02  |  7.33e+00 +/- 3.2e-02  |
 test_solve_fixed_iter                   |     +0.36 +/- 58.21    | +1.78e-02 +/- 2.86e+00 |  4.93e+00 +/- 2.1e+00  |  4.91e+00 +/- 1.9e+00  |

[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

[dpanici]

Make a blocked grad and hess to allow rev mode with this objective (or have it use FD for grad and hess)

[dpanici]

Add a basic tutorial (call subprocess etc but just to run a simple python script not an actual external code)

[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