Add center=True option for CoilSetMinDistance and PlasmaCoilSetMinDistance

[dpanici]

For small shaping coils, the center of the coil can be used as a proxy for the position along the coil, and one can use this in the distance objectives instead of the points along the coil. This can be much cheaper than using positions along the coil if the number of coils is large.

This PR adds center flag to CoilSetMinDistance and PlasmaCoilSetMinDistance which will use the center of the coil to compute distances from each coil.

Waiting on @dpanici to look into segmented sum to make this work so that we can have some coils use centers and others use the full coil

[github-actions[bot]]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -0.62 +/- 1.94     | -3.52e-03 +/- 1.10e-02 |  5.65e-01 +/- 7.8e-03  |  5.68e-01 +/- 7.7e-03  |
 test_equilibrium_init_medres            |     -0.69 +/- 4.23     | -3.14e-02 +/- 1.91e-01 |  4.49e+00 +/- 1.4e-01  |  4.52e+00 +/- 1.3e-01  |
 test_equilibrium_init_highres           |     -0.91 +/- 3.35     | -4.65e-02 +/- 1.71e-01 |  5.07e+00 +/- 1.4e-01  |  5.12e+00 +/- 1.0e-01  |
 test_objective_compile_dshape_current   |     -0.50 +/- 1.34     | -1.76e-02 +/- 4.77e-02 |  3.53e+00 +/- 2.6e-02  |  3.55e+00 +/- 4.0e-02  |
 test_objective_compute_dshape_current   |     -0.08 +/- 3.04     | -2.78e-06 +/- 1.05e-04 |  3.44e-03 +/- 8.3e-05  |  3.45e-03 +/- 6.4e-05  |
 test_objective_jac_dshape_current       |     +3.01 +/- 15.11    | +9.35e-04 +/- 4.69e-03 |  3.20e-02 +/- 3.6e-03  |  3.11e-02 +/- 3.0e-03  |
 test_perturb_2                          |     -1.02 +/- 2.46     | -1.82e-01 +/- 4.38e-01 |  1.76e+01 +/- 3.8e-01  |  1.78e+01 +/- 2.1e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.24 +/- 0.74     | +3.59e-02 +/- 1.12e-01 |  1.51e+01 +/- 9.6e-02  |  1.51e+01 +/- 5.7e-02  |
 test_proximal_freeb_jac                 |     +0.69 +/- 9.62     | +3.28e-02 +/- 4.60e-01 |  4.82e+00 +/- 3.6e-01  |  4.78e+00 +/- 2.9e-01  |
 test_solve_fixed_iter_compiled          |     +0.62 +/- 3.46     | +1.10e-01 +/- 6.11e-01 |  1.78e+01 +/- 4.2e-01  |  1.77e+01 +/- 4.4e-01  |
 test_LinearConstraintProjection_build   |     -0.29 +/- 2.36     | -2.52e-02 +/- 2.02e-01 |  8.55e+00 +/- 1.2e-01  |  8.57e+00 +/- 1.6e-01  |
 test_objective_compute_ripple_spline    |     +1.57 +/- 4.02     | +4.92e-03 +/- 1.26e-02 |  3.19e-01 +/- 5.7e-03  |  3.14e-01 +/- 1.1e-02  |
 test_objective_grad_ripple_spline       |     -0.12 +/- 2.93     | -1.47e-03 +/- 3.63e-02 |  1.24e+00 +/- 2.5e-02  |  1.24e+00 +/- 2.7e-02  |
 test_build_transform_fft_midres         |     -0.30 +/- 3.09     | -2.23e-03 +/- 2.32e-02 |  7.49e-01 +/- 1.9e-02  |  7.51e-01 +/- 1.3e-02  |
 test_build_transform_fft_highres        |     -0.88 +/- 2.65     | -9.01e-03 +/- 2.72e-02 |  1.02e+00 +/- 2.0e-02  |  1.03e+00 +/- 1.9e-02  |
 test_equilibrium_init_lowres            |     -0.75 +/- 3.44     | -3.40e-02 +/- 1.56e-01 |  4.51e+00 +/- 1.0e-01  |  4.54e+00 +/- 1.2e-01  |
 test_objective_compile_atf              |     +0.30 +/- 1.74     | +1.97e-02 +/- 1.15e-01 |  6.64e+00 +/- 8.0e-02  |  6.62e+00 +/- 8.2e-02  |
 test_objective_compute_atf              |     -3.12 +/- 16.87    | -2.83e-04 +/- 1.53e-03 |  8.80e-03 +/- 9.6e-04  |  9.08e-03 +/- 1.2e-03  |
 test_objective_jac_atf                  |     +1.39 +/- 4.10     | +2.18e-02 +/- 6.43e-02 |  1.59e+00 +/- 5.5e-02  |  1.57e+00 +/- 3.3e-02  |
 test_perturb_1                          |     -0.19 +/- 1.83     | -2.95e-02 +/- 2.79e-01 |  1.52e+01 +/- 1.5e-01  |  1.52e+01 +/- 2.3e-01  |
 test_proximal_jac_atf                   |     -0.37 +/- 1.06     | -2.81e-02 +/- 8.15e-02 |  7.68e+00 +/- 5.8e-02  |  7.71e+00 +/- 5.8e-02  |
 test_proximal_freeb_compute             |     -0.86 +/- 2.93     | -1.58e-03 +/- 5.39e-03 |  1.83e-01 +/- 3.2e-03  |  1.84e-01 +/- 4.3e-03  |
 test_solve_fixed_iter                   |     -2.65 +/- 1.43     | -8.45e-01 +/- 4.56e-01 |  3.11e+01 +/- 3.5e-01  |  3.19e+01 +/- 2.9e-01  |
 test_objective_compute_ripple           |     -0.14 +/- 1.10     | -3.92e-03 +/- 3.05e-02 |  2.76e+00 +/- 2.5e-02  |  2.76e+00 +/- 1.7e-02  |
 test_objective_grad_ripple              |     +1.10 +/- 1.99     | +5.48e-02 +/- 9.92e-02 |  5.05e+00 +/- 6.0e-02  |  5.00e+00 +/- 7.9e-02  |

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 95.75%. Comparing base (ef16bad) to head (bfe1ccc).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1689      +/-   ##
==========================================
- Coverage   95.76%   95.75%   -0.01%     
==========================================
  Files         101      101              
  Lines       26616    26631      +15     
==========================================
+ Hits        25488    25501      +13     
- Misses       1128     1130       +2     

Files with missing lines	Coverage Δ
desc/coils.py	98.05% <100.00%> (+0.03%)	:arrow_up:
desc/objectives/_coils.py	99.38% <100.00%> (+<0.01%)	:arrow_up:

... and 1 file with indirect coverage changes

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dpanici]

This would be most useful if we could allow some of the coils to be treated fully while others are treated with their centers... this may make it a bit messier though

[dpanici]

segmented sum JAX

[dpanici]

segmented sum JAX

see https://github.com/PlasmaControl/DESC/blob/b20e730f1dc932afd51cf193aa2017fa07fefa32/desc/objectives/_reconstruction.py#L247 for example use

[unalmis]

segmented sum JAX

see https://github.com/PlasmaControl/DESC/blob/b20e730f1dc932afd51cf193aa2017fa07fefa32/desc/objectives/_reconstruction.py#L247 for example use

last year i remember all the jax segment ops were significantly less efficient (both memory and speed) than constructing the sparse matrix and multiplying

[dpanici]

segmented sum JAX

see https://github.com/PlasmaControl/DESC/blob/b20e730f1dc932afd51cf193aa2017fa07fefa32/desc/objectives/_reconstruction.py#L247 for example use

last year i remember all the jax segment ops were significantly less efficient (both memory and speed) than constructing the sparse matrix and multiplying

At least for me in this case, the ops were easy to use and way better than a for loop. I will think about making a sparse matrix but if I can get it working with ops I am fine with that for now.

By sparse matrix you mean enforcing the segmented sum by just doing a big matrix multiply and masking out the entries you don't want?