DESC icon indicating copy to clipboard operation
DESC copied to clipboard

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

Open dpanici opened this issue 1 year ago • 65 comments

  • Adds function run_regcoil to _current_potential.py that implements the REGCOIL algorithm (Landreman, (2017)) for surface current normal field optimization
    • Can specify current_helicity to determine if resulting contours correspond to helical topology (current_helicity not equal to 0) or modular (current_helicityequal to 0)
    • if multiple values of the regularization parameter are input, will return a family of surface current fields (as a list) corresponding to the solution at each regularization value
    • 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|^2, the regularization term from surface current in the REGCOIL algorithm, with w being the objective weight which act as 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.
  • Adds a tutorial showing these features

Resolves #578

Remaining TODO (Or for future PRs):

  • [ ] debug result when M=0, N=0 mode is in FourierCurrentPotentialField basis
  • [ ] 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)
  • [ ] 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)
  • [x] 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)
  • [x] add area element to run_regcoil
  • [ ] change run_regcoil name
  • [x] 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
  • [x] address notebook comments

Old TODO

TODO list:

  • [x] make more efficient by only evaluating Bn on one field period

  • [x] 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)

  • [x] simplify function

  • [x] 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"

  • [x] replace biot function with one present in DESC already, or add it to a util function

  • [x] add tests

  • [x] add MagneticField class for CurrentPotentialField (so can use in finding Bnorm, field line tracing, etc) #592 -> merge this and things will be way more efficient

  • [x] 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)

  • [x] Add tests for coil finding with helicity ratio positive

  • [x] add tests for coil finding with |helicity ratio| >1

  • [x] 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)

  • [x] objective to do this (one for quad flux and one for regularizaiton)

  • [x] method of surfac current field to cut coils

  • [x] make sure source grid is sym=False

  • [x] find helical coils in smarter way, not using full 0->2pi in NFP

  • [x] finite beta (just need to include the plasma contribution) - might wait on #725 for this just to use the singular integral stuff, should add warnings for now about if used on a non-vacuum equilibrium

  • [x] replace contour finding matplotlib call with scikit contour finding function (to avoid having to make a figure when creating coils)

  • [x] Add tutorial for using these functions for coil optimization (Once #627 is implemented, use that algorithm in the tutorial)

  • [x] change field trace in tutorial to trace at multiple zeta, not just zeta=0

  • [x] fix tutorial regcoil calls to use new functions

  • [x] 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

  • [x] impement coil cutting field period symmetry for modular coils

CLEANUP TODO FOR WORKFLOW:

  • [x] clean up field line tracing from current potential script (once #592 is in, dont need a separate script)
  • [x] clean up field line tracing from Coilset (again the necessity of a script is questionable)
  • [x] 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)
  • [x] separate functions for run regcoil and coil cutting into private functions that return coilset or families of surface current fields
  • [x] remove plotting from the run regcoil and coil cutting
  • [x] add plot util for the regcoil plots
  • [x] add compute function for min dist between coils in a coilset
  • [x] check if order of things matters for quad flux objective (or remove in favor of the one in master)
  • [x] weight Bn by area (surf jacobian)

dpanici avatar Jul 13 '23 22:07 dpanici