FourierPlanarCurve basis option

[ddudt]

Adds an option for the center and normal vectors of FourierPlanarCurve to be specified in $(R,\phi,Z)$ coordinates, instead of the previous limitation to Cartesian coordinates.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.96%. Comparing base (7418b99) to head (0f9a6fe).
Report is 1423 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1017   +/-   ##
=======================================
  Coverage   94.95%   94.96%           
=======================================
  Files          87       87           
  Lines       21716    21740   +24     
=======================================
+ Hits        20621    20645   +24     
  Misses       1095     1095           

Files with missing lines	Coverage Δ
desc/coils.py	96.55% <100.00%> (ø)
desc/compute/_curve.py	99.45% <100.00%> (+0.03%)	⬆️
desc/geometry/core.py	95.91% <ø> (ø)
desc/geometry/curve.py	93.25% <100.00%> (+0.06%)	⬆️

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +1.87 +/- 4.93     | +9.38e-03 +/- 2.47e-02 |  5.10e-01 +/- 2.2e-02  |  5.01e-01 +/- 1.0e-02  |
 test_build_transform_fft_midres         |     +0.92 +/- 2.45     | +5.29e-03 +/- 1.41e-02 |  5.82e-01 +/- 1.0e-02  |  5.77e-01 +/- 9.8e-03  |
 test_build_transform_fft_highres        |     +0.84 +/- 1.41     | +8.12e-03 +/- 1.37e-02 |  9.78e-01 +/- 7.2e-03  |  9.70e-01 +/- 1.2e-02  |
 test_equilibrium_init_lowres            |     +1.91 +/- 3.17     | +7.27e-02 +/- 1.20e-01 |  3.87e+00 +/- 6.9e-02  |  3.80e+00 +/- 9.8e-02  |
 test_equilibrium_init_medres            |     +1.35 +/- 2.63     | +5.80e-02 +/- 1.13e-01 |  4.34e+00 +/- 7.7e-02  |  4.29e+00 +/- 8.2e-02  |
 test_equilibrium_init_highres           |     +1.98 +/- 1.91     | +1.18e-01 +/- 1.14e-01 |  6.10e+00 +/- 7.3e-02  |  5.98e+00 +/- 8.7e-02  |
 test_objective_compile_dshape_current   |     +0.34 +/- 1.09     | +1.19e-02 +/- 3.87e-02 |  3.55e+00 +/- 1.3e-02  |  3.53e+00 +/- 3.6e-02  |
 test_objective_compile_atf              |     +1.07 +/- 1.34     | +7.41e-02 +/- 9.33e-02 |  7.02e+00 +/- 7.7e-02  |  6.94e+00 +/- 5.2e-02  |
 test_objective_compute_dshape_current   |     +1.78 +/- 2.19     | +7.22e-05 +/- 8.89e-05 |  4.13e-03 +/- 7.1e-05  |  4.06e-03 +/- 5.3e-05  |
 test_objective_compute_atf              |     +0.94 +/- 2.15     | +1.75e-04 +/- 4.01e-04 |  1.88e-02 +/- 2.3e-04  |  1.86e-02 +/- 3.3e-04  |
 test_objective_jac_dshape_current       |     -1.02 +/- 5.34     | -4.14e-04 +/- 2.17e-03 |  4.01e-02 +/- 1.7e-03  |  4.06e-02 +/- 1.4e-03  |
 test_objective_jac_atf                  |     -0.63 +/- 1.54     | -1.17e-02 +/- 2.85e-02 |  1.84e+00 +/- 2.3e-02  |  1.86e+00 +/- 1.7e-02  |
 test_perturb_1                          |     -0.73 +/- 3.81     | -9.35e-02 +/- 4.85e-01 |  1.27e+01 +/- 2.9e-01  |  1.28e+01 +/- 3.9e-01  |
 test_perturb_2                          |     +0.80 +/- 2.47     | +1.40e-01 +/- 4.33e-01 |  1.76e+01 +/- 2.8e-01  |  1.75e+01 +/- 3.3e-01  |
 test_proximal_jac_atf                   |     +0.45 +/- 1.04     | +3.18e-02 +/- 7.39e-02 |  7.17e+00 +/- 5.0e-02  |  7.14e+00 +/- 5.4e-02  |
 test_proximal_freeb_compute             |     +1.14 +/- 0.89     | +2.19e-03 +/- 1.70e-03 |  1.94e-01 +/- 1.2e-03  |  1.92e-01 +/- 1.3e-03  |
 test_proximal_freeb_jac                 |     +0.39 +/- 1.21     | +2.85e-02 +/- 8.88e-02 |  7.39e+00 +/- 7.2e-02  |  7.36e+00 +/- 5.3e-02  |

[dpanici]

I feel like there is a better way to do this using a property of FourierPlanarCurve called basis_in and then having it be obtained through transforms in the compute method, instead of needing to redefine the compute method for this one curve.

[dpanici]

Try this in a simple optimization (does not have to be meaningful or involve an equilibrium, just optimize the coil to change the curvature or something), @kianorr has encountered issues with JAX when a string is present in transforms and so I wonder if that occurs here too

[dpanici]

@kianorr I added a test that I thought would trigger the JAX issue you had, but I got another error, can you help debug?

    c = FourierPlanarCurve()
    objective = ObjectiveFunction(CoilLength(c, target=11))
    optimizer = Optimizer("fmintr")
    (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
    np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)

has an error when the objective is built

desc/optimize/optimizer.py:217: in optimize
    objective.build(verbose=verbose)
desc/objectives/objective_funs.py:173: in build
    objective.build(use_jit=self.use_jit, verbose=verbose)
desc/objectives/_coils.py:337: in build
    super().build(use_jit=use_jit, verbose=verbose)
desc/objectives/_coils.py:153: in build
    self._grid = tree_map(
../../miniconda3/envs/desc-env/lib/python3.12/site-packages/jax/_src/tree_util.py:312: in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
../../miniconda3/envs/desc-env/lib/python3.12/site-packages/jax/_src/tree_util.py:312: in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

x = array([0., 1., 0.])

    lambda x: LinearGrid(
>       N=2 * x.N + 5, NFP=getattr(x, "NFP", 1), endpoint=False
    ),
    self.things[0],
    is_leaf=lambda x: is_single_coil(x),
E   AttributeError: 'numpy.ndarray' object has no attribute 'N'

desc/objectives/_coils.py:155: AttributeError

[dpanici]

Ah it seems that the coil objectives cannot handle curves being passed. That is unfortunate, I will make a separate issue. The latest test I added shows the JAX issue though, sorry @ddudt seems my suggestion was a bad one as I had forgotten about this issue. The kwarg method you had originally might be the way to go

[ddudt]

Ah it seems that the coil objectives cannot handle curves being passed. That is unfortunate, I will make a separate issue. The latest test I added shows the JAX issue though, sorry @ddudt seems my suggestion was a bad one as I had forgotten about this issue. The kwarg method you had originally might be the way to go

See this test in PR #977 that uses this new planar curve parameterization in a coil optimization: https://github.com/PlasmaControl/DESC/blob/dp/coil-coil-dist/tests/test_examples.py#L1301

[dpanici]

Ah it seems that the coil objectives cannot handle curves being passed. That is unfortunate, I will make a separate issue. The latest test I added shows the JAX issue though, sorry @ddudt seems my suggestion was a bad one as I had forgotten about this issue. The kwarg method you had originally might be the way to go

See this test in PR #977 that uses this new planar curve parameterization in a coil optimization: https://github.com/PlasmaControl/DESC/blob/dp/coil-coil-dist/tests/test_examples.py#L1301

The test I added here still fails on that PR (both when only the single coil is optimized, or if it is included in a MixedCoilSet). I think the way I suggested to do this was wrong and we might have to use your original way

[ddudt]

Ah it seems that the coil objectives cannot handle curves being passed. That is unfortunate, I will make a separate issue. The latest test I added shows the JAX issue though, sorry @ddudt seems my suggestion was a bad one as I had forgotten about this issue. The kwarg method you had originally might be the way to go

See this test in PR #977 that uses this new planar curve parameterization in a coil optimization: https://github.com/PlasmaControl/DESC/blob/dp/coil-coil-dist/tests/test_examples.py#L1301

The test I added here still fails on that PR (both when only the single coil is optimized, or if it is included in a MixedCoilSet). I think the way I suggested to do this was wrong and we might have to use your original way

But the way this "basis" option is implemented at the moment is equivalent to the "method" option for computing the length of a SplineXYZCurve. If a SplineXYZCoil works with the CoilLength objective, then I think what we are trying to do here should also work.

[dpanici]

Ah it seems that the coil objectives cannot handle curves being passed. That is unfortunate, I will make a separate issue. The latest test I added shows the JAX issue though, sorry @ddudt seems my suggestion was a bad one as I had forgotten about this issue. The kwarg method you had originally might be the way to go

See this test in PR #977 that uses this new planar curve parameterization in a coil optimization: https://github.com/PlasmaControl/DESC/blob/dp/coil-coil-dist/tests/test_examples.py#L1301

The test I added here still fails on that PR (both when only the single coil is optimized, or if it is included in a MixedCoilSet). I think the way I suggested to do this was wrong and we might have to use your original way

But the way this "basis" option is implemented at the moment is equivalent to the "method" option for computing the length of a SplineXYZCurve. If a SplineXYZCoil works with the CoilLength objective, then I think what we are trying to do here should also work.

SplineXYZCoil also do not work it seems, they throw the same JAX error, I just checked

[ddudt]

Ah it seems that the coil objectives cannot handle curves being passed. That is unfortunate, I will make a separate issue. The latest test I added shows the JAX issue though, sorry @ddudt seems my suggestion was a bad one as I had forgotten about this issue. The kwarg method you had originally might be the way to go

See this test in PR #977 that uses this new planar curve parameterization in a coil optimization: https://github.com/PlasmaControl/DESC/blob/dp/coil-coil-dist/tests/test_examples.py#L1301

The test I added here still fails on that PR (both when only the single coil is optimized, or if it is included in a MixedCoilSet). I think the way I suggested to do this was wrong and we might have to use your original way

But the way this "basis" option is implemented at the moment is equivalent to the "method" option for computing the length of a SplineXYZCurve. If a SplineXYZCoil works with the CoilLength objective, then I think what we are trying to do here should also work.

SplineXYZCoil also do not work it seems, they throw the same JAX error, I just checked

@dpanici I reverted back to my original solution.

@kianorr you can copy this solution for your spline method in #970. Also note that you can replace the new example tests in this PR when you resolve #1021.