Add Coil Arclength Differential Variance Objective

[dpanici]

• Adds objective similar to the one detailed in Wechsung 2021 to penalize the variance of the arclength differential along a curve, which was claimed to be useful for avoiding issues with parameterization non-uniqueness in optimization.
• add a test for this
• Add simplified version, this is equal to the simsopt implementation with nintervals="full"

From reference:

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     +0.97 +/- 8.86     | +6.02e-03 +/- 5.49e-02 |  6.26e-01 +/- 4.0e-02  |  6.20e-01 +/- 3.8e-02  |
 test_build_transform_fft_highres        |     -1.74 +/- 3.83     | -1.78e-02 +/- 3.91e-02 |  1.00e+00 +/- 2.0e-02  |  1.02e+00 +/- 3.4e-02  |
 test_equilibrium_init_lowres            |     -3.78 +/- 5.26     | -1.51e-01 +/- 2.11e-01 |  3.85e+00 +/- 1.7e-01  |  4.00e+00 +/- 1.2e-01  |
 test_objective_compile_atf              |     -1.91 +/- 3.99     | -1.56e-01 +/- 3.25e-01 |  7.98e+00 +/- 2.2e-01  |  8.13e+00 +/- 2.4e-01  |
 test_objective_compute_atf              |     -6.33 +/- 3.79     | -6.97e-04 +/- 4.17e-04 |  1.03e-02 +/- 1.7e-04  |  1.10e-02 +/- 3.8e-04  |
 test_objective_jac_atf                  |     -2.33 +/- 2.65     | -4.69e-02 +/- 5.32e-02 |  1.97e+00 +/- 3.2e-02  |  2.01e+00 +/- 4.2e-02  |
 test_perturb_1                          |     -7.10 +/- 2.73     | -9.71e-01 +/- 3.73e-01 |  1.27e+01 +/- 2.5e-01  |  1.37e+01 +/- 2.8e-01  |
 test_proximal_jac_atf                   |     -1.27 +/- 1.08     | -1.05e-01 +/- 8.97e-02 |  8.18e+00 +/- 3.8e-02  |  8.29e+00 +/- 8.1e-02  |
 test_proximal_freeb_compute             |     -0.84 +/- 1.04     | -1.58e-03 +/- 1.94e-03 |  1.85e-01 +/- 1.5e-03  |  1.87e-01 +/- 1.2e-03  |
 test_build_transform_fft_lowres         |     -5.33 +/- 12.00    | -3.51e-02 +/- 7.90e-02 |  6.23e-01 +/- 3.1e-02  |  6.58e-01 +/- 7.3e-02  |
 test_equilibrium_init_medres            |     -0.49 +/- 9.20     | -2.43e-02 +/- 4.61e-01 |  4.99e+00 +/- 2.1e-01  |  5.01e+00 +/- 4.1e-01  |
 test_equilibrium_init_highres           |     -1.73 +/- 9.62     | -1.08e-01 +/- 6.00e-01 |  6.13e+00 +/- 4.4e-01  |  6.23e+00 +/- 4.1e-01  |
 test_objective_compile_dshape_current   |     -9.94 +/- 9.83     | -4.54e-01 +/- 4.50e-01 |  4.12e+00 +/- 2.0e-01  |  4.57e+00 +/- 4.0e-01  |
 test_objective_compute_dshape_current   |    -19.99 +/- 7.29     | -9.43e-04 +/- 3.44e-04 |  3.77e-03 +/- 1.1e-04  |  4.72e-03 +/- 3.3e-04  |
 test_objective_jac_dshape_current       |    -17.72 +/- 14.19    | -9.98e-03 +/- 7.99e-03 |  4.63e-02 +/- 4.1e-03  |  5.63e-02 +/- 6.9e-03  |
 test_perturb_2                          |    -11.30 +/- 7.88     | -2.57e+00 +/- 1.79e+00 |  2.02e+01 +/- 1.4e+00  |  2.28e+01 +/- 1.1e+00  |
 test_proximal_freeb_jac                 |    -14.99 +/- 20.75    | -1.43e+00 +/- 1.99e+00 |  8.14e+00 +/- 7.9e-01  |  9.57e+00 +/- 1.8e+00  |
 test_solve_fixed_iter                   |    -13.02 +/- 58.89    | -7.96e-01 +/- 3.60e+00 |  5.32e+00 +/- 2.4e+00  |  6.11e+00 +/- 2.7e+00  |

[codecov]

Codecov Report

Attention: Patch coverage is 96.66667% with 1 line in your changes missing coverage. Please review.

Project coverage is 95.48%. Comparing base (e5d31c2) to head (8e0c13c).

Files with missing lines	Patch %	Lines
desc/objectives/_coils.py	96.55%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1113   +/-   ##
=======================================
  Coverage   95.48%   95.48%           
=======================================
  Files          96       96           
  Lines       23965    23995   +30     
=======================================
+ Hits        22882    22912   +30     
  Misses       1083     1083           

Files with missing lines	Coverage Δ
desc/compute/_curve.py	100.00% <ø> (ø)
desc/objectives/__init__.py	100.00% <ø> (ø)
desc/transform.py	93.29% <100.00%> (+0.01%)	⬆️
desc/objectives/_coils.py	99.29% <96.55%> (+0.05%)	⬆️

... and 1 file with indirect coverage changes

[dpanici]

check for which curves this makes sense to have

[f0uriest]

why is this needed?

If you want to skip the checks the better way is just to call _compute directly, which is what we do in the objectives

[dpanici]

CoilSet objects have logic to deal with computing its tree structure, which would be annoying to implement at the objective level every time in order to get to the base Coil object that can then be called with _compute. Having a jitable kwarg here to skip the unneeded checks lets us use coilset.compute inside of objectives.

This is due to an uncaught bug in CoilObjectives, any objective that requires compute quantitiy that needs a transform with derivatives to be built (like "x" "x_s"), the has_transforms check will fail due to a JAX error on us calling to_list on the derivatives. I can post the exact error later

[f0uriest]

OK, for an immediate fix it sounds like it might be better to just mark Transform._derivs as a static attribute?

Bigger picture though it would be nice to have the option to ignore checks when calling compute, this is why we call _compute in all the other objectives but that's kind of hacky. We could move all these checks to their own function, and then compute would be something like

def compute(..., check=True):
    if check:
        _check_inputs(...)
    return _compute(...)

Then we can use compute everywhere safely, just sometimes needing to pass in check=False

[f0uriest]

This is due to an uncaught bug in CoilObjectives, any objective that requires compute quantitiy that needs a transform with derivatives to be built (like "x" "x_s"), the has_transforms check will fail due to a JAX error on us calling to_list on the derivatives. I can post the exact error later

I'm confused by this, don't objectives like CoilCurvature and CoilTorsion require derivatives of the position? Why do they not fail?

[dpanici]

I also was confused, I think it is somehow bc the target compute quantitiy itself has no transforms/deriv requirements, only the dependencies. I did not dig much deeper into it though