Add Vector Potential Calculation to Coil classes and Most MagneticField Classes

[dpanici]

Add utilities for calculating the magnetic vector potential for fields and coils/currents (assuming the Coulomb gauge). These will be useful for reconstruction as the most efficient way to calculate flux through a diagnostic flux loop is to integrate the vector potential. The utilities here should allow for calculating the vector potential from most sources of interest: coils and the plasma contribution (through K=nxB and the general biot-savart law here, as the loops are located outside from the plasma)

• Add general biot savart law for vector potential
• Add Hanson-Hirhsman biot savart law for coil vector potential
• Add quadrature biot savart law for coil vector potential
• add to current potential
• add tests for current potential (tests loop integral against enclosed toroidal flux)
• add to coils
• add test for coils for simple loop on z-axis
• add a test against enclosed toroidal flux doing a loop integral (like the one in this paper perhaps)
  • understand why flipping orientation of the coil parameterization does not flip sign of A? Check signs of A and x_s, though even just the A does not change sign as I expect... but I'd expect xs to change sign when the parameterizwion is flipped
• add to the analytic fields
• add tests for the analytic fields
• Add VectorPotentialField
• Add to SplineMagneticField (from mgrid)
• note all gauges (try to put Coulomb Gauge for all) TF and VF are, and the rest are either from AD or mgrid, so we can't really check that
• Use Vector potential for ToroidalFlux objective (possibly as an option, since not all B fields have a vector potential field)

Resolves #1023

[codecov]

Codecov Report

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

Project coverage is 95.44%. Comparing base (d2e9a2c) to head (98053b3).
Report is 72 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/magnetic_fields/_core.py	97.05%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1026      +/-   ##
==========================================
+ Coverage   95.41%   95.44%   +0.02%     
==========================================
  Files          95       95              
  Lines       23186    23406     +220     
==========================================
+ Hits        22124    22340     +216     
- Misses       1062     1066       +4     

Files with missing lines	Coverage Δ
desc/coils.py	97.46% <100.00%> (+0.21%)	⬆️
desc/magnetic_fields/__init__.py	100.00% <ø> (ø)
desc/magnetic_fields/_current_potential.py	99.44% <100.00%> (+0.03%)	⬆️
desc/objectives/_coils.py	99.17% <100.00%> (+0.03%)	⬆️
desc/magnetic_fields/_core.py	96.56% <97.05%> (+0.16%)	⬆️

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.27 +/- 10.81    | +1.49e-03 +/- 5.90e-02 |  5.47e-01 +/- 5.4e-02  |  5.46e-01 +/- 2.5e-02  |
 test_equilibrium_init_medres            |     -0.99 +/- 4.81     | -4.27e-02 +/- 2.08e-01 |  4.29e+00 +/- 2.0e-01  |  4.33e+00 +/- 5.2e-02  |
 test_equilibrium_init_highres           |     +1.52 +/- 1.92     | +8.65e-02 +/- 1.10e-01 |  5.79e+00 +/- 7.8e-02  |  5.71e+00 +/- 7.7e-02  |
 test_objective_compile_dshape_current   |     +0.62 +/- 1.14     | +2.45e-02 +/- 4.47e-02 |  3.94e+00 +/- 3.7e-02  |  3.92e+00 +/- 2.5e-02  |
 test_objective_compute_dshape_current   |     -0.40 +/- 3.89     | -1.43e-05 +/- 1.38e-04 |  3.53e-03 +/- 8.2e-05  |  3.54e-03 +/- 1.1e-04  |
 test_objective_jac_dshape_current       |     +2.73 +/- 8.20     | +1.13e-03 +/- 3.40e-03 |  4.26e-02 +/- 3.1e-03  |  4.15e-02 +/- 1.5e-03  |
 test_perturb_2                          |     +0.75 +/- 1.27     | +1.36e-01 +/- 2.28e-01 |  1.82e+01 +/- 1.5e-01  |  1.80e+01 +/- 1.7e-01  |
 test_proximal_freeb_jac                 |     +1.00 +/- 1.60     | +7.50e-02 +/- 1.20e-01 |  7.61e+00 +/- 8.3e-02  |  7.53e+00 +/- 8.7e-02  |
 test_solve_fixed_iter                   |     +0.98 +/- 60.04    | +4.91e-02 +/- 3.01e+00 |  5.06e+00 +/- 2.1e+00  |  5.01e+00 +/- 2.2e+00  |

[f0uriest]

div/grad/curl in jax (assumes cartesian coordinates, will need some extra stuff to work in cylindrical)

def grad(fun, argnum):
    return jax.jacfwd(fun, argnum)

def div(fun, argnum):
    def divf(*args, **kwargs):
        return jnp.trace(jax.jacfwd(fun, argnum)(*args, **kwargs))
    return divf

def curl(fun, argnum):
    def curlf(*args, **kwargs):
        J = jax.jacfwd(fun, argnum)(*args, **kwargs)
        return jnp.array([J[2,1]-J[1,2],
                         J[0,2]-J[2,0],
                         J[1,0]-J[0,1]])
    return curlf

[ddudt]

My main comment is that I think for nearly all classes the code would be much simpler if you get rid of the _compute_A_or_B methods. The one exception might be SplineMagneticField where there is actually some shared logic.

[ddudt]

This method seems pointless if it is always called by either compute_magnetic_field or compute_magnetic_vector_potential

[ddudt]

See comments on the subclasses, but if this is getting superseded by child classes then it should be an abstract method. Same for compute_magnetic_field and compute_magnetic_vector_potential

[f0uriest]

the point is to avoid duplicated code. A lot of the preprocessing is the same, the only thing that really changes is the integrand

[ddudt]

There are only about a dozen lines of duplicate code, and that is mostly just for checking "rpz" vs "xyz" basis which we already do in all the compute functions too. I would rather make a single util function to do that check and conversion that we can call everywhere. Because the header/docstring of this function alone is like 40 lines of code that also gets duplicated many times, even if it isn't "real code".

[dpanici]

for SplineMagneticField and CurrentPotentialField the duplicated code is 100+ lines

I don't have a strong preference either way but I do think this current way is cleaner for avoiding duplicated code, and this is all under the hood stuff anyways. Unless there is a stronger case for splitting this up I want to keep it as it is now.

[ddudt]

I mentioned in my other comment that this is fine for those classes where there is actually a lot of duplicated code. My point is that for most of the classes, the only duplicated code is 10 lines of stuff related to converting between rpz vs xyz basis. And in those cases I don't think it makes sense to add 40+ lines of code to every class for something that can be a simple util function call.

[ddudt]

In other words, move the actual duplicate code to a separate util function, and then all the code that is specific to A vs B can be moved to their respective compute functions.

[ddudt]

Same comment as I already mentioned for the Coil classes: I think it will be fewer lines of code and easier to understand if you get rid of these methods and just put the logic for either A or B in their respective methods.

[dpanici]

remove _ fxn for toroidal field and vertical field and poloidal field