Add OmnigenityHarmonics Objective and new OmnigenousField types

[DMCXE]

Description

This pull request introduces the OOPS (Omnigenity OPtimization like quasi-Symmetry) method into DESC, as described in [1]. OOPS leverages the concept of a homeomorphism to enable the optimization of omnigenity, as well as other concepts like pseudo-symetry and piecewise omnigenity, in spectral space, analogous to the optimization of quasisymmetry. A key advantage of this method is that it is independent of any specific parameterization. One can use the Cary-Shasharina mapping [2], the Landreman-Catto mapping [3], the current DESC mapping, or any others. It only requires a method for constructing the target magnetic field. Generally, it has better roubustness. The goal of this PR is to implement OOPS within DESC's flexible and efficient framework.
To facilitate discussion, I have created this Draft Pull Request. See same issue in #1846

Core Changes

The main changes included in this PR are:

• New OmnigenousField Types: Two new parameterizations for omnigenous fields have been added in magnetic_field/_core.py:
  • OmnigenousFieldOOPS: The example mapping used in the OOPS paper [1].
  • OmnigenousFieldLCForm: The mapping form proposed by Landreman et al. [3], referred to here as the "LCForm".
• New OmnigenityHarmonics Objective: A new objective, OmnigenityHarmonics, has been added to objectives/_omnigenity.py. This objective enables the direct optimization of OmnigenousField, OmnigenousFieldOOPS, and OmnigenousFieldLCForm in spectral space.
• New Computations: Necessary parameters have been registered in compute/_omnigenity.py. The core addition is the function _B_omni_nonsymmetric, which calculates the non-symmetric components of the omnigenous field, B(η,α).
• New Plotting Functions: To support visualization of these new field types, plot_boozer_surface_XXX routines for OmnigenousFieldOOPS and OmnigenousFieldLCForm have been implemented in plotting.py.

Examples

Here are some examples in Ref. [1]. They can be optimized using DESC. We have implemented an optimization workflow in DESC with objectives, procedures, and resolutions that are nearly identical to the corresponding SIMSOPT scripts. Following the same process, starting from a circular torus, we use a resolution continuation method to optimize the configuration up to a spectral resolution of M=N=7. Depending on whether jax-cache is enabled, a precisely poloidal omnigenous (PO) configuration can be optimized within 10-15 minutes on a single NVIDIA RTX 5090 GPU. (With a configuration nearly identical to the one described in Ref. [1])

Known I/O Limitation

While the I/O interface can successfully save and load the degrees of freedom of OmnigenousFieldLCForm (s_dofs, d_dofs), it does not currently support serializing the function definitions themselves. This means that when loading from a saved state, the user must manually re-supply the Python functions for S_function and D_function.

References

[1] H. Liu, G. Yu, et al., "Optimizing omnigenity like quasisymmetry for stellarators," arxiv:2502.09350 (2025).
[2] J. R. Cary and S. G. Shasharina, Helical Plasma Confinement Devices with Good Confinement Properties, Phys. Rev. Lett. 78, 674 (1997).
[3] M. Landreman and P. J. Catto, "Omnigenity as generalized quasisymmetry," Physics of Plasmas 19, 056103 (2012).

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    1.81 %    |     3.912e+03      |     3.983e+03      |    70.78     |       38.06        |       34.75        |
  test_proximal_jac_w7x_with_eq_update   |    1.48 %    |     6.512e+03      |     6.608e+03      |    96.38     |       155.35       |       158.39       |
  test_proximal_freeb_jac                |    0.51 %    |     1.318e+04      |     1.325e+04      |    67.17     |       80.21        |       80.21        |
  test_proximal_freeb_jac_blocked        |    0.96 %    |     7.610e+03      |     7.683e+03      |    73.39     |       72.81        |       72.60        |
  test_proximal_freeb_jac_batched        |    0.03 %    |     7.619e+03      |     7.621e+03      |     2.52     |       73.42        |       70.98        |
  test_proximal_jac_ripple               |   -2.02 %    |     3.493e+03      |     3.422e+03      |    -70.68    |       59.97        |       60.17        |
  test_proximal_jac_ripple_bounce1d      |   -1.32 %    |     3.532e+03      |     3.485e+03      |    -46.68    |       77.35        |       76.65        |
  test_eq_solve                          |    2.93 %    |     1.984e+03      |     2.042e+03      |    58.14     |       128.29       |       126.75       |

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

[dpanici]

For the tests to pass you should add the new objectives to the "specials" list as they need extra arguments compared to the typical objective, for example see

DESC/tests/test_objective_funs.py

Line 3169 in 4e87d19

Omnigenity,

If you can do this we can then make sure the normal tests pass. You should also add some tests like what exist for Omnigenity in test_objective_funs.py

[rahulgaur104]

Another big code-related change would be to create a PR from a branch on the DESC repo, not on a fork from your repo.
I don't remember how but it helps with tests and merging into master.

Edit: Ah, I remember some of the issues. Checking out and updating your PR becomes difficult on a cluster.
I also vaguely remember some problems with the tests.

[dpanici]

some TODO items I see as needed for this PR:

• Update the objective functions test as I had pointed out, so that the tests pass
• add tests for the new functionality
  • in test_objective_funs.py, covering the different use cases

    • DESC/tests/test_objective_funs.py

      Line 3522 in c704244

      def test_compute_scalar_resolution_omnigenity(self):

    • DESC/tests/test_objective_funs.py

      Line 3935 in c704244

      def test_objective_no_nangrad_omnigenity(self, helicity):

  • adding an example(s) in test_examples

    • DESC/tests/test_examples.py

      Line 953 in c704244

      def test_omnigenity_qa():

    • DESC/tests/test_examples.py

      Line 1036 in c704244

      def test_omnigenity_optimization():

    • DESC/tests/test_examples.py

      Line 1128 in c704244

      def test_omnigenity_proximal():

• see if you can merge the different _compute variable names (LCform, OOPS, etc) for things like D_list into the just one declaration with multiple parametrizations listed for it (I think for D_list, S_list you can do this for sure)
• add comments in the _compute to mark which equations in the paper different things correspond to, this will help with clarity
• expand the existing plot_boozer_surface to be able to plot the new field types, instead of making an entirely new plotting function which I think re-uses a lot of the same code
• Perhaps add a tutorial highlighting this, maybe add as another example in the existing Omnigenity tutorial notebook?

For the point on S_function and D_function, perhaps a possible solution could be to use pickle to save them as .p files? I assume the functional forms are more complex than just fourier series (and hence must be stored instead of just storing the coefficients) for the piecewise omnigenity equilibria specifically? otherwise d_n and s_n would be sufficient info?

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

❌ Patch coverage is 97.57463% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.81%. Comparing base (71d59bb) to head (faa8396).

Files with missing lines	Patch %	Lines
desc/magnetic_fields/_core.py	96.80%	6 Missing ⚠️
desc/objectives/_omnigenity.py	97.79%	4 Missing ⚠️
desc/compute/_omnigenity.py	98.07%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1845      +/-   ##
==========================================
+ Coverage   95.78%   95.81%   +0.03%     
==========================================
  Files         101      101              
  Lines       27728    28258     +530     
==========================================
+ Hits        26558    27076     +518     
- Misses       1170     1182      +12     

Files with missing lines	Coverage Δ
desc/compute/data_index.py	95.89% <ø> (ø)
desc/magnetic_fields/__init__.py	100.00% <ø> (ø)
desc/objectives/__init__.py	100.00% <ø> (ø)
desc/plotting.py	95.49% <100.00%> (+0.02%)	⬆️
desc/compute/_omnigenity.py	99.26% <98.07%> (-0.74%)	⬇️
desc/objectives/_omnigenity.py	97.33% <97.79%> (+0.26%)	⬆️
desc/magnetic_fields/_core.py	96.53% <96.80%> (+0.06%)	⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

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