Option to use proper normalization for QS objectives

[YigitElma]

During stage-1 optimization, if the volume or the average magnetic field is not constrained, the strength of the magnetic field can be reduced by inflating the equilibrium (while still having the same Psi). The actual definition of these objectives (for example, in Landreman&Paul 2022) already includes these normalizations. Our precomputed normalizations cannot prevent these kinds of "cheats" in some cases.

This PR adds the option to choose which equation to use by the new mode argument. The default keeps the old behavior.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -0.31 %    |     4.239e+03      |     4.225e+03      |    -13.22    |       25.52        |       22.55        |
  test_proximal_jac_w7x_with_eq_update   |   -0.52 %    |     6.856e+03      |     6.821e+03      |    -35.69    |       118.48       |       118.82       |
  test_proximal_freeb_jac                |   -0.09 %    |     1.354e+04      |     1.353e+04      |    -11.77    |       63.26        |       65.99        |
  test_proximal_freeb_jac_blocked        |   -0.09 %    |     7.876e+03      |     7.869e+03      |    -6.77     |       57.12        |       55.90        |
  test_proximal_freeb_jac_batched        |   -0.20 %    |     7.871e+03      |     7.856e+03      |    -15.46    |       55.20        |       55.10        |
  test_proximal_jac_ripple               |   -0.85 %    |     3.850e+03      |     3.818e+03      |    -32.71    |       40.96        |       40.59        |
  test_proximal_jac_ripple_bounce1d      |   -1.82 %    |     3.970e+03      |     3.898e+03      |    -72.20    |       51.16        |       50.85        |
  test_eq_solve                          |   -0.35 %    |     1.833e+03      |     1.827e+03      |    -6.47     |       40.31        |       40.25        |
  test_objective_quadratic_flux_jac      |   -0.35 %    |     1.897e+03      |     1.890e+03      |    -6.69     |       27.70        |       27.59        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.36%. Comparing base (187e98e) to head (4d10e24).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2304      +/-   ##
==========================================
+ Coverage   94.34%   94.36%   +0.01%     
==========================================
  Files         101      101              
  Lines       29042    29073      +31     
==========================================
+ Hits        27400    27435      +35     
+ Misses       1642     1638       -4     

Files with missing lines	Coverage Δ
desc/objectives/_omnigenity.py	97.61% <100.00%> (+0.56%)	⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +2.42 +/- 3.62     | +2.03e-02 +/- 3.04e-02 |  8.58e-01 +/- 2.4e-02  |  8.38e-01 +/- 1.8e-02  |
 test_equilibrium_init_lowres            |     +0.02 +/- 4.04     | +1.05e-03 +/- 2.58e-01 |  6.39e+00 +/- 2.1e-01  |  6.39e+00 +/- 1.5e-01  |
 test_objective_compile_atf              |     +4.10 +/- 4.47     | +2.52e-01 +/- 2.74e-01 |  6.39e+00 +/- 2.2e-01  |  6.14e+00 +/- 1.7e-01  |
 test_objective_compute_atf              |     +7.65 +/- 18.40    | +1.57e-04 +/- 3.76e-04 |  2.20e-03 +/- 2.0e-04  |  2.05e-03 +/- 3.2e-04  |
 test_objective_jac_atf                  |     +4.87 +/- 3.73     | +7.62e-02 +/- 5.83e-02 |  1.64e+00 +/- 5.0e-02  |  1.56e+00 +/- 3.0e-02  |
 test_perturb_1                          |     +1.31 +/- 3.19     | +1.54e-01 +/- 3.75e-01 |  1.19e+01 +/- 3.6e-01  |  1.18e+01 +/- 1.0e-01  |
 test_proximal_jac_atf                   |     +1.10 +/- 2.82     | +5.78e-02 +/- 1.48e-01 |  5.31e+00 +/- 1.1e-01  |  5.26e+00 +/- 1.0e-01  |
 test_proximal_freeb_compute             |     +3.96 +/- 3.27     | +6.54e-03 +/- 5.41e-03 |  1.72e-01 +/- 3.8e-03  |  1.65e-01 +/- 3.8e-03  |
 test_solve_fixed_iter                   |     -0.77 +/- 4.99     | -1.97e-01 +/- 1.27e+00 |  2.52e+01 +/- 1.1e+00  |  2.54e+01 +/- 5.8e-01  |
 test_LinearConstraintProjection_build   |     -6.44 +/- 4.59     | -4.72e-01 +/- 3.37e-01 |  6.86e+00 +/- 2.4e-01  |  7.34e+00 +/- 2.4e-01  |
 test_objective_compute_ripple           |     -0.69 +/- 4.56     | -1.64e-03 +/- 1.09e-02 |  2.36e-01 +/- 4.1e-03  |  2.38e-01 +/- 1.0e-02  |
 test_objective_grad_ripple              |     -1.51 +/- 3.81     | -1.36e-02 +/- 3.44e-02 |  8.88e-01 +/- 1.5e-02  |  9.02e-01 +/- 3.1e-02  |
 test_objective_quadratic_flux_compute   |     +0.98 +/- 1.23     | +2.09e-04 +/- 2.62e-04 |  2.15e-02 +/- 1.9e-04  |  2.13e-02 +/- 1.8e-04  |
 test_build_transform_fft_midres         |     +0.18 +/- 2.98     | +1.52e-03 +/- 2.54e-02 |  8.53e-01 +/- 9.1e-03  |  8.52e-01 +/- 2.4e-02  |
 test_build_transform_fft_highres        |     +0.56 +/- 3.80     | +6.36e-03 +/- 4.31e-02 |  1.14e+00 +/- 3.9e-02  |  1.14e+00 +/- 1.8e-02  |
 test_equilibrium_init_medres            |     -3.55 +/- 6.43     | -2.48e-01 +/- 4.49e-01 |  6.74e+00 +/- 3.1e-01  |  6.98e+00 +/- 3.2e-01  |
 test_objective_compile_dshape_current   |     -3.57 +/- 3.15     | -1.44e-01 +/- 1.27e-01 |  3.89e+00 +/- 4.5e-02  |  4.04e+00 +/- 1.2e-01  |
 test_objective_compute_dshape_current   |     -2.60 +/- 15.53    | -1.61e-05 +/- 9.61e-05 |  6.03e-04 +/- 4.4e-05  |  6.19e-04 +/- 8.5e-05  |
 test_objective_jac_dshape_current       |     -0.76 +/- 27.99    | -1.89e-04 +/- 6.93e-03 |  2.46e-02 +/- 4.0e-03  |  2.48e-02 +/- 5.6e-03  |
 test_perturb_2                          |     -3.48 +/- 4.78     | -5.35e-01 +/- 7.35e-01 |  1.48e+01 +/- 3.2e-01  |  1.54e+01 +/- 6.6e-01  |
 test_proximal_jac_atf_with_eq_update    |     -0.58 +/- 1.24     | -7.20e-02 +/- 1.54e-01 |  1.23e+01 +/- 1.3e-01  |  1.24e+01 +/- 8.3e-02  |
 test_proximal_freeb_jac                 |     +0.42 +/- 3.50     | +1.91e-02 +/- 1.61e-01 |  4.61e+00 +/- 9.2e-02  |  4.59e+00 +/- 1.3e-01  |
 test_solve_fixed_iter_compiled          |     +0.72 +/- 2.72     | +4.32e-02 +/- 1.62e-01 |  6.01e+00 +/- 1.5e-01  |  5.97e+00 +/- 6.4e-02  |
 test_objective_compute_ripple_bounce1d  |     -1.05 +/- 3.77     | -2.99e-03 +/- 1.07e-02 |  2.82e-01 +/- 4.0e-03  |  2.85e-01 +/- 1.0e-02  |
 test_objective_grad_ripple_bounce1d     |     +0.58 +/- 4.14     | +5.60e-03 +/- 4.00e-02 |  9.71e-01 +/- 1.6e-02  |  9.66e-01 +/- 3.6e-02  |
 test_objective_quadratic_flux_jac       |     +1.74 +/- 1.74     | +3.84e-02 +/- 3.85e-02 |  2.25e+00 +/- 3.5e-02  |  2.21e+00 +/- 1.7e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.

[dpanici]

I'd prefer to just silently ignore normalize=True when the hatted metric is requested, as technically it IS normalized. Just set the normalization in the build to 1 in those cases

[YigitElma]

Also add the equations in the docstring