Add JAX implementation of REGCOIL algorithm and Add Ability to Discretize Current Potentials into Coilsets

[dpanici]

Question: It would be nice to update our coil tutorial to have a section where we use this functionality to intiialize the coils for optimization, can I add this in this PR?

• Adds function solve_regularized_surface_current to desc.magnetic_fields module that implements the REGCOIL algorithm (Landreman, (2017)) for surface current normal field optimization
  • Can specify the tuple current_helicity=(q,p) to determine if resulting contours correspond to helical topology (both (q,p) not equal to 0), modular (p equal to 0 and q nonzero), PF (p nonzero but q zero, so coils don't link the plasma poloidally but they do close toroidally like a PF coil for a tokamak), or windowpane/saddle (p and q both zero). q is number of poloidal transits a current potentiail contour (and consequently, a coil) makes per field period, while p is the number of toroidal transits per field period (not really true toroidal transits, but think of it as how many field periods it takes for a coil to return to the same poloidal position on surface, the larger the p is for fixed q, the more PF-coil-like the coil topology becomes)
  • Can also specify regularization_type which can be simple or regcoil, the difference explained here, tldr is simple is not guaranteed to monotonically decrease $\chi^2_K$ as regularization increases, but gives the same qualitative behavior and is less expensive to run.
• Adds method To_CoilSet to FourierCurrentPotentialField which implements a coil cutting algorithm to discretize the surface current into coils
  • works for both modular and helical coils
  • uses skimage.measure.find_contours instead of matplotlib's deprecated contour finding algorithm, so this PR adds a new dependency in skimage.
• Adds a new objective SurfaceCurrentRegularization (which minimizes w*|K|, the regularization term from surface current in the REGCOIL algorithm, with w being the objective weight which act as the sqrt of the regularization parameter)
  • use of both this and the QuadraticFlux objective allows for REGCOIL solutions to be obtained through the optimization framework, and combined with other objectives as well.
  • QuadraticFlux has a sqrt_area_weighting flag added to make the area weighting be the sqrt of the surface Jacobian, which makes it more in line with the REGCOIL algorithm (as the cost squared and summed should be equal to $\chi^2_B$). Defaults to False which is the current behavior on master
• Adds a new tutorial showing how to useREGCOIL features.

Resolves #578

Remaining TODO (Or for future PRs):

• debug result when M=0, N=0 mode is in FourierCurrentPotentialField basis (i.e. when sym_phi="cos" or False) -> add warning if wrong phi symmetry or something
• add options to scan over external_TF fraction and over helicity_ratio (that way only need to find Jacobian of Bn_SV(phi_mn) once) (maybe just add option to return Jacobian? then can make function way simpler, and have an example of how to use to do a scan...) -> removed this as it was a bit specific to NT-TAO and made the code a bit messy, though I think it could be a reasonable thing to have as once you have the jacobian, it is free to change the external field... maybe adding a kwarg that uses a passed-in jacobian could be a workaround that keeps this feature
• easier way to set current potential helicity when not doing the run_regcoil method - maybe a flag like (set_G_automatically which sets it based off of the integrated B_zeta from the eq and the external field, if provided) - this is done in the tutorial, it is simple enough so I will leave it as is
• Maybe attempt to speedup by not using JAX for jacobian but instead constructing manually, compare speeds (maybe an improvement for later...)
• use plot_2d inside plot_regcoil_outputs for Phi and K (but not Bn, as it would be expensive to recompute Bplasma every time we call plot_2d)
• implement checks for coil cutting stell symmetry for modular coils (a bit more annoying as the resulting coilset will not be good if there is a coil that crosses the zeta=0 plane... could just check if any contours cross zeta=0 or zeta=pi/nfp and throw an error if so)
• add area element to run_regcoil
• change run_regcoil name
• change heliicty to p/q to match my paper
• use cholesky if single alpha is passed instead of SVD
• use ax instead of plt in plotting util
• address notebook comments

Old TODO

TODO list:

• make more efficient by only evaluating Bn on one field period
• address aliasing issues (need large grids for both eval and source to accurately describe Bn for even a simple surface current density - This is mainly solved by increasing source_grid_N, which must be sufficiently large (something like NFP * basis_N * 2 or *3 at least, while source_grid_M can be 1.5 or 2 times basis_M. eval_grids can be 1.5 or 2 times basis_M and basis_N)
• simplify function
• rename helicity_ratio to something more clear, possibly re-define current potential G sign so that a positive helicity ratio corresponds to a positive slope of the contours of constant current potential? or at least explain better what it is, maybe just name it "coil helicity"
• replace biot function with one present in DESC already, or add it to a util function
• add tests
• add MagneticField class for CurrentPotentialField (so can use in finding Bnorm, field line tracing, etc) Dp/current potential field #592 -> merge this and things will be way more efficient
• add util function for finding Bnorm from a coilset on a surface possibly? could be a function of the MagneticField class which takes in the Surface object (tho maybe need to specify which surface since need to compute n_rho for toroidal surface, but probabnly not important to be compatible with Poincare surfaces right now)
• Add tests for coil finding with helicity ratio positive
• add tests for coil finding with |helicity ratio| >1
• in coil finding, check which direction the contour is heading, and ensure that sign of current is correct (can get K direction from n x grad(phi), dot with first contour pt - last contour point , use sign of that as indication of if current should be flipped or not)
• objective to do this (one for quad flux and one for regularizaiton)
• method of surfac current field to cut coils
• make sure source grid is sym=False
• find helical coils in smarter way, not using full 0->2pi in NFP
• finite beta (just need to include the plasma contribution) - might wait on Add free boundary equilibrium objectives #725 for this just to use the singular integral stuff, should add warnings for now about if used on a non-vacuum equilibrium
• replace contour finding matplotlib call with scikit contour finding function (to avoid having to make a figure when creating coils)
• Add tutorial for using these functions for coil optimization (Once Constant Offset Surface Algorithm #627 is implemented, use that algorithm in the tutorial)
• change field trace in tutorial to trace at multiple zeta, not just zeta=0
• fix tutorial regcoil calls to use new functions
• Use precise QA for regcoil test, as it does not take that long at lowres and is a better test than the barely nonaxisymmetric ellipse I am using currently
• impement coil cutting field period symmetry for modular coils

CLEANUP TODO FOR WORKFLOW:

• clean up field line tracing from current potential script (once Dp/current potential field #592 is in, dont need a separate script)
• clean up field line tracing from Coilset (again the necessity of a script is questionable)
• clean up coilset BNORM script (necessity of a separate script is questionable, can put in one script with calculate_Bnorm method of MagneticField, possibly with a plot option)
• separate functions for run regcoil and coil cutting into private functions that return coilset or families of surface current fields
• remove plotting from the run regcoil and coil cutting
• add plot util for the regcoil plots
• add compute function for min dist between coils in a coilset
• check if order of things matters for quad flux objective (or remove in favor of the one in master)
• weight Bn by area (surf jacobian)

[dpanici]

Ah this actually is branched off of #494 , should I point this PR towards there, or just wait to merge until that one is ready and merged in?

[codecov]

Codecov Report

Attention: Patch coverage is 98.70801% with 5 lines in your changes missing coverage. Please review.

Project coverage is 95.58%. Comparing base (d62cddd) to head (59cb549).

Files with missing lines	Patch %	Lines
desc/magnetic_fields/_current_potential.py	99.36%	2 Missing ⚠️
desc/objectives/_coils.py	96.29%	2 Missing ⚠️
desc/plotting.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #579      +/-   ##
==========================================
+ Coverage   95.52%   95.58%   +0.06%     
==========================================
  Files          96       96              
  Lines       24005    24385     +380     
==========================================
+ Hits        22930    23308     +378     
- Misses       1075     1077       +2     

Files with missing lines	Coverage Δ
desc/compute/_surface.py	100.00% <100.00%> (ø)
desc/magnetic_fields/__init__.py	100.00% <100.00%> (ø)
desc/magnetic_fields/_core.py	96.61% <ø> (ø)
desc/objectives/__init__.py	100.00% <ø> (ø)
desc/plotting.py	96.16% <83.33%> (-0.07%)	⬇️
desc/magnetic_fields/_current_potential.py	99.39% <99.36%> (-0.06%)	⬇️
desc/objectives/_coils.py	99.16% <96.29%> (-0.37%)	⬇️

... and 4 files with indirect coverage changes

[github-actions]

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

flake8

desc/field_line_tracing_DESC_from_coilset.py|11 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|12 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|12 col 1| 'desc.magnetic_fields.ToroidalMagneticField' imported but unused
desc/field_line_tracing_DESC_from_coilset.py|12 col 1| 'desc.magnetic_fields.SumMagneticField' imported but unused
desc/field_line_tracing_DESC_from_coilset.py|12 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|12 col 89| line too long (93 > 88 characters)
desc/field_line_tracing_DESC_from_coilset.py|13 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|13 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|14 col 1| 'desc.io.load' imported but unused
desc/field_line_tracing_DESC_from_coilset.py|14 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|15 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|15 col 1| 'scipy.interpolate.interp1d' imported but unused
desc/field_line_tracing_DESC_from_coilset.py|15 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|16 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|17 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|17 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|18 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|18 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|19 col 1| isort found an unexpected blank line in imports
desc/field_line_tracing_DESC_from_coilset.py|20 col 1| isort found an unexpected blank line in imports
desc/field_line_tracing_DESC_from_coilset.py|21 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|21 col 1| 'desc.coils.Coil' imported but unused
desc/field_line_tracing_DESC_from_coilset.py|21 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|22 col 1| isort found an unexpected blank line in imports
desc/field_line_tracing_DESC_from_coilset.py|23 col 1| isort found an unexpected missing import
desc/field_line_tracing_DESC_from_coilset.py|23 col 1| isort found an unexpected missing import
desc/field_line_tracing_DESC_from_coilset.py|23 col 1| isort found an unexpected missing import
desc/field_line_tracing_DESC_from_coilset.py|23 col 1| isort found an unexpected missing import
desc/field_line_tracing_DESC_from_coilset.py|23 col 1| isort found an unexpected missing import
desc/field_line_tracing_DESC_from_coilset.py|23 col 1| isort found an unexpected missing import
desc/field_line_tracing_DESC_from_coilset.py|31 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|31 col 1| 'utils.get_coilset_from_MAKEGRID_or_regcoil' imported but unused
desc/field_line_tracing_DESC_from_coilset.py|31 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|33 col 21| trailing whitespace
desc/field_line_tracing_DESC_from_coilset.py|35 col 1| isort found an import in the wrong position
desc/field_line_tracing_DESC_from_coilset.py|35 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|40 col 23| trailing whitespace
desc/field_line_tracing_DESC_from_coilset.py|43 col 89| line too long (95 > 88 characters)
desc/field_line_tracing_DESC_from_coilset.py|46 col 14| trailing whitespace
desc/field_line_tracing_DESC_from_coilset.py|69 col 1| redefinition of unused 'os' from line 18
desc/field_line_tracing_DESC_from_coilset.py|69 col 1| module level import not at top of file
desc/field_line_tracing_DESC_from_coilset.py|70 col 1| isort expected 1 blank line in imports, found 0
desc/field_line_tracing_DESC_from_coilset.py|73 col 89| line too long (91 > 88 characters)
desc/field_line_tracing_DESC_from_coilset.py|74 col 89| line too long (91 > 88 characters)
desc/field_line_tracing_DESC_from_coilset.py|89 col 1| blank line contains whitespace
desc/field_line_tracing_DESC_from_coilset.py|115 col 1| too many blank lines (3)
desc/field_line_tracing_DESC_from_coilset.py|115 col 89| line too long (122 > 88 characters)
desc/field_line_tracing_DESC_from_coilset.py|117 col 1| blank line at end of file

[dpanici]

closing until files are cleaned up

[dpanici]

reopening so tests will run

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -0.81 +/- 1.41     | -1.54e-04 +/- 2.68e-04 |  1.88e-02 +/- 2.0e-04  |  1.90e-02 +/- 1.8e-04  |
 test_build_transform_fft_midres         |     +0.92 +/- 1.23     | +1.02e-03 +/- 1.37e-03 |  1.12e-01 +/- 9.9e-04  |  1.11e-01 +/- 9.5e-04  |
 test_build_transform_fft_highres        |     +0.51 +/- 0.80     | +2.65e-03 +/- 4.14e-03 |  5.18e-01 +/- 1.8e-03  |  5.15e-01 +/- 3.7e-03  |
 test_equilibrium_init_lowres            |     -0.56 +/- 1.71     | -2.45e-03 +/- 7.52e-03 |  4.36e-01 +/- 6.4e-03  |  4.39e-01 +/- 4.0e-03  |
 test_equilibrium_init_medres            |     -0.60 +/- 1.81     | -3.39e-03 +/- 1.02e-02 |  5.60e-01 +/- 4.6e-03  |  5.64e-01 +/- 9.1e-03  |
 test_equilibrium_init_highres           |     -0.60 +/- 1.14     | -5.94e-03 +/- 1.13e-02 |  9.86e-01 +/- 7.6e-03  |  9.92e-01 +/- 8.4e-03  |
 test_objective_compile_dshape_current   |     -1.65 +/- 7.48     | -8.70e-02 +/- 3.95e-01 |  5.19e+00 +/- 2.6e-01  |  5.28e+00 +/- 2.9e-01  |
 test_objective_compile_atf              |     +0.62 +/- 5.38     | +1.01e-01 +/- 8.71e-01 |  1.63e+01 +/- 7.0e-01  |  1.62e+01 +/- 5.1e-01  |
 test_objective_compute_dshape_current   |     -0.68 +/- 2.09     | -2.25e-05 +/- 6.96e-05 |  3.31e-03 +/- 6.7e-05  |  3.33e-03 +/- 1.9e-05  |
 test_objective_compute_atf              |     -1.96 +/- 0.90     | -2.22e-04 +/- 1.02e-04 |  1.11e-02 +/- 6.6e-05  |  1.13e-02 +/- 7.8e-05  |
 test_objective_jac_dshape_current       |     +2.19 +/- 13.98    | +2.81e-03 +/- 1.79e-02 |  1.31e-01 +/- 1.6e-02  |  1.28e-01 +/- 7.4e-03  |
 test_objective_jac_atf                  |     +2.01 +/- 3.54     | +1.52e-01 +/- 2.68e-01 |  7.71e+00 +/- 1.4e-01  |  7.56e+00 +/- 2.3e-01  |
 test_perturb_1                          |     -1.07 +/- 16.01    | -9.26e-02 +/- 1.39e+00 |  8.58e+00 +/- 9.9e-01  |  8.68e+00 +/- 9.8e-01  |
 test_perturb_2                          |     -0.35 +/- 5.53     | -5.78e-02 +/- 9.05e-01 |  1.63e+01 +/- 6.1e-01  |  1.64e+01 +/- 6.7e-01  |

[f0uriest]

still not crazy about this name