Adding a piecewise omnigenity objective

[rahulgaur104]

This PR add the piecewise omnigenity objective as described by Velasco during the Simons meeting.

This will allows us to do much faster an memory efficient optimizations while maintaining approximate omnigenity and maybe use reverse-mode differentiation.

The expression for the magnetic field is given in a screenshot of Velasco's Simons slide.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +3.55 +/- 5.53     | +1.94e-02 +/- 3.03e-02 |  5.67e-01 +/- 2.5e-02  |  5.47e-01 +/- 1.7e-02  |
 test_build_transform_fft_midres         |     +3.23 +/- 5.11     | +2.04e-02 +/- 3.23e-02 |  6.51e-01 +/- 2.9e-02  |  6.31e-01 +/- 1.3e-02  |
 test_build_transform_fft_highres        |     +0.11 +/- 3.10     | +1.13e-03 +/- 3.15e-02 |  1.02e+00 +/- 1.7e-02  |  1.02e+00 +/- 2.7e-02  |
 test_equilibrium_init_lowres            |     -4.39 +/- 4.04     | -1.87e-01 +/- 1.72e-01 |  4.08e+00 +/- 9.3e-02  |  4.26e+00 +/- 1.4e-01  |
 test_equilibrium_init_medres            |     -2.15 +/- 4.34     | -9.79e-02 +/- 1.97e-01 |  4.45e+00 +/- 1.2e-01  |  4.55e+00 +/- 1.5e-01  |
 test_equilibrium_init_highres           |     -2.59 +/- 3.65     | -1.54e-01 +/- 2.17e-01 |  5.80e+00 +/- 1.0e-01  |  5.96e+00 +/- 1.9e-01  |
 test_objective_compile_dshape_current   |     +0.94 +/- 2.56     | +3.75e-02 +/- 1.02e-01 |  4.03e+00 +/- 8.8e-02  |  3.99e+00 +/- 5.3e-02  |
 test_objective_compile_atf              |     +3.29 +/- 1.79     | +2.66e-01 +/- 1.45e-01 |  8.35e+00 +/- 1.2e-01  |  8.09e+00 +/- 7.9e-02  |
 test_objective_compute_dshape_current   |     +9.73 +/- 6.29     | +3.39e-04 +/- 2.19e-04 |  3.82e-03 +/- 1.8e-04  |  3.48e-03 +/- 1.3e-04  |
 test_objective_compute_atf              |     +0.05 +/- 4.22     | +5.90e-06 +/- 4.80e-04 |  1.14e-02 +/- 3.7e-04  |  1.14e-02 +/- 3.1e-04  |
 test_objective_jac_dshape_current       |     +1.72 +/- 7.30     | +7.28e-04 +/- 3.09e-03 |  4.30e-02 +/- 2.8e-03  |  4.23e-02 +/- 1.2e-03  |
 test_objective_jac_atf                  |     +1.65 +/- 2.65     | +3.20e-02 +/- 5.16e-02 |  1.98e+00 +/- 4.6e-02  |  1.94e+00 +/- 2.3e-02  |
 test_perturb_1                          |     +3.18 +/- 1.62     | +4.21e-01 +/- 2.14e-01 |  1.36e+01 +/- 1.3e-01  |  1.32e+01 +/- 1.7e-01  |
 test_perturb_2                          |     +1.23 +/- 3.49     | +2.26e-01 +/- 6.41e-01 |  1.86e+01 +/- 2.7e-01  |  1.84e+01 +/- 5.8e-01  |
 test_proximal_jac_atf                   |     +1.09 +/- 1.38     | +8.93e-02 +/- 1.14e-01 |  8.32e+00 +/- 6.5e-02  |  8.23e+00 +/- 9.4e-02  |
 test_proximal_freeb_compute             |     +1.70 +/- 1.56     | +3.13e-03 +/- 2.88e-03 |  1.88e-01 +/- 2.5e-03  |  1.84e-01 +/- 1.4e-03  |
 test_proximal_freeb_jac                 |     +1.94 +/- 3.06     | +1.49e-01 +/- 2.35e-01 |  7.83e+00 +/- 2.2e-01  |  7.68e+00 +/- 7.2e-02  |
 test_solve_fixed_iter                   |     +0.77 +/- 61.23    | +3.90e-02 +/- 3.12e+00 |  5.13e+00 +/- 2.3e+00  |  5.09e+00 +/- 2.1e+00  |

[rahulgaur104]

When I try to import the PiecewiseOmnigenity objective, I get the following error:

DESC version 0.11.1+555.g6474d1e7b,using JAX backend, jax version=0.4.24, jaxlib version=0.4.24, dtype=float64
Using device: CPU, with 23.07 GB available memory
Traceback (most recent call last):
  File "/home/rgaur/pwO_import_test.py", line 2, in <module>
    from desc.objectives import PiecewiseOmnigenity
  File "/home/rgaur/DESC/desc/objectives/__init__.py", line 3, in <module>
    from ._bootstrap import BootstrapRedlConsistency
  File "/home/rgaur/DESC/desc/objectives/_bootstrap.py", line 6, in <module>
    from desc.compute import get_profiles, get_transforms
  File "/home/rgaur/DESC/desc/compute/__init__.py", line 31, in <module>
    from . import (
  File "/home/rgaur/DESC/desc/compute/_omnigenity.py", line 507, in <module>
    def _B_piecewise_omni(params, transforms, profiles, data, **kwargs):
  File "/home/rgaur/DESC/desc/compute/data_index.py", line 170, in _decorator
    raise ValueError(
ValueError: Can't register function with unknown parameterization: desc.magnetic_fields._core.PiecewiseOmnigenousField

[f0uriest]

When I try to import the PiecewiseOmnigenity objective, I get the following error:

DESC version 0.11.1+555.g6474d1e7b,using JAX backend, jax version=0.4.24, jaxlib version=0.4.24, dtype=float64
Using device: CPU, with 23.07 GB available memory
Traceback (most recent call last):
  File "/home/rgaur/pwO_import_test.py", line 2, in <module>
    from desc.objectives import PiecewiseOmnigenity
  File "/home/rgaur/DESC/desc/objectives/__init__.py", line 3, in <module>
    from ._bootstrap import BootstrapRedlConsistency
  File "/home/rgaur/DESC/desc/objectives/_bootstrap.py", line 6, in <module>
    from desc.compute import get_profiles, get_transforms
  File "/home/rgaur/DESC/desc/compute/__init__.py", line 31, in <module>
    from . import (
  File "/home/rgaur/DESC/desc/compute/_omnigenity.py", line 507, in <module>
    def _B_piecewise_omni(params, transforms, profiles, data, **kwargs):
  File "/home/rgaur/DESC/desc/compute/data_index.py", line 170, in _decorator
    raise ValueError(
ValueError: Can't register function with unknown parameterization: desc.magnetic_fields._core.PiecewiseOmnigenousField

You need to add the new class to desc.compute.data_index._class_inheritance