plot elongation

[ddudt]

Resolves #804.

Does the same grid resampling that we do for functions of $\rho$, but now also for $\zeta$.

The grid_type data index tag is no longer used, but I keep it because we should start using this in the future. This PR resolves an issue about plotting functions of $\zeta$ vs $\rho$, but that tag will still be useful for determining the proper type of grid to use (e.g. "quad" vs "linear").

I added a new test and also changed the other 1D plotting tests to get examples so they don't depend on doing an equilibrium solve.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +3.94 +/- 8.15     | +4.73e-04 +/- 9.77e-04 |  1.25e-02 +/- 7.0e-04  |  1.20e-02 +/- 6.8e-04  |
 test_build_transform_fft_midres         |     +1.19 +/- 6.43     | +1.04e-03 +/- 5.60e-03 |  8.81e-02 +/- 3.5e-03  |  8.70e-02 +/- 4.4e-03  |
 test_build_transform_fft_highres        |     +3.43 +/- 2.86     | +1.50e-02 +/- 1.25e-02 |  4.53e-01 +/- 7.0e-03  |  4.38e-01 +/- 1.0e-02  |
 test_equilibrium_init_lowres            |     +2.43 +/- 2.46     | +1.85e-02 +/- 1.87e-02 |  7.79e-01 +/- 9.7e-03  |  7.60e-01 +/- 1.6e-02  |
 test_equilibrium_init_medres            |     -0.42 +/- 1.87     | -5.76e-03 +/- 2.56e-02 |  1.36e+00 +/- 2.1e-02  |  1.37e+00 +/- 1.5e-02  |
 test_equilibrium_init_highres           |     -0.54 +/- 1.41     | -2.19e-02 +/- 5.71e-02 |  4.02e+00 +/- 4.2e-02  |  4.04e+00 +/- 3.8e-02  |
 test_objective_compile_dshape_current   |     -1.53 +/- 7.51     | -6.48e-02 +/- 3.18e-01 |  4.17e+00 +/- 2.2e-01  |  4.23e+00 +/- 2.3e-01  |
 test_objective_compile_atf              |     +0.34 +/- 7.19     | +3.05e-02 +/- 6.39e-01 |  8.92e+00 +/- 6.2e-01  |  8.89e+00 +/- 1.7e-01  |
 test_objective_compute_dshape_current   |     +0.27 +/- 7.24     | +5.72e-06 +/- 1.52e-04 |  2.11e-03 +/- 1.5e-04  |  2.10e-03 +/- 3.0e-05  |
+test_objective_compute_atf              |     -6.76 +/- 2.19     | -5.09e-04 +/- 1.65e-04 |  7.02e-03 +/- 1.5e-04  |  7.53e-03 +/- 7.3e-05  |
 test_objective_jac_dshape_current       |     -4.01 +/- 8.98     | -1.86e-03 +/- 4.18e-03 |  4.47e-02 +/- 3.7e-03  |  4.65e-02 +/- 2.0e-03  |
 test_objective_jac_atf                  |     -0.22 +/- 6.86     | -4.80e-03 +/- 1.47e-01 |  2.14e+00 +/- 8.4e-02  |  2.14e+00 +/- 1.2e-01  |
 test_perturb_1                          |     +0.41 +/- 13.72    | +3.40e-02 +/- 1.14e+00 |  8.32e+00 +/- 7.8e-01  |  8.28e+00 +/- 8.3e-01  |
 test_perturb_2                          |     -0.15 +/- 5.57     | -2.15e-02 +/- 7.81e-01 |  1.40e+01 +/- 6.4e-01  |  1.40e+01 +/- 4.5e-01  |

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +3.09 +/- 5.28     | +3.87e-04 +/- 6.63e-04 |  1.29e-02 +/- 6.5e-04  |  1.26e-02 +/- 1.1e-04  |
 test_build_transform_fft_midres         |     -1.06 +/- 1.23     | -9.98e-04 +/- 1.16e-03 |  9.30e-02 +/- 1.0e-03  |  9.40e-02 +/- 5.6e-04  |
 test_build_transform_fft_highres        |     -1.85 +/- 1.01     | -8.68e-03 +/- 4.75e-03 |  4.62e-01 +/- 3.5e-03  |  4.70e-01 +/- 3.2e-03  |
 test_equilibrium_init_lowres            |     -0.60 +/- 1.54     | -4.70e-03 +/- 1.22e-02 |  7.84e-01 +/- 9.5e-03  |  7.88e-01 +/- 7.5e-03  |
 test_equilibrium_init_medres            |     +0.32 +/- 1.59     | +4.48e-03 +/- 2.25e-02 |  1.42e+00 +/- 1.6e-02  |  1.42e+00 +/- 1.5e-02  |
 test_equilibrium_init_highres           |     -1.36 +/- 1.02     | -5.75e-02 +/- 4.29e-02 |  4.15e+00 +/- 3.8e-02  |  4.21e+00 +/- 2.1e-02  |
 test_objective_compile_dshape_current   |     -1.38 +/- 8.49     | -6.05e-02 +/- 3.72e-01 |  4.32e+00 +/- 2.6e-01  |  4.38e+00 +/- 2.7e-01  |
 test_objective_compile_atf              |     -0.89 +/- 6.26     | -8.14e-02 +/- 5.74e-01 |  9.09e+00 +/- 3.6e-01  |  9.18e+00 +/- 4.5e-01  |
 test_objective_compute_dshape_current   |     +1.05 +/- 2.53     | +2.22e-05 +/- 5.34e-05 |  2.13e-03 +/- 4.6e-05  |  2.11e-03 +/- 2.8e-05  |
 test_objective_compute_atf              |     -1.01 +/- 1.34     | -7.62e-05 +/- 1.01e-04 |  7.47e-03 +/- 7.9e-05  |  7.55e-03 +/- 6.3e-05  |
 test_objective_jac_dshape_current       |     -0.39 +/- 12.96    | -1.81e-04 +/- 6.00e-03 |  4.61e-02 +/- 3.2e-03  |  4.63e-02 +/- 5.1e-03  |
 test_objective_jac_atf                  |     +0.39 +/- 5.25     | +8.37e-03 +/- 1.14e-01 |  2.18e+00 +/- 6.3e-02  |  2.17e+00 +/- 9.5e-02  |
 test_perturb_1                          |     +0.16 +/- 13.19    | +1.34e-02 +/- 1.13e+00 |  8.59e+00 +/- 8.8e-01  |  8.58e+00 +/- 7.1e-01  |
 test_perturb_2                          |     +3.06 +/- 5.75     | +4.46e-01 +/- 8.38e-01 |  1.50e+01 +/- 5.8e-01  |  1.46e+01 +/- 6.0e-01  |

[ddudt]

@unalmis I need your help to make the computation of "a_major/a_minor" work on symmetric grids

[unalmis]

@unalmis I need your help to make the computation of "a_major/a_minor" work on symmetric grids

I will work on this now

[unalmis]

@ddudt The computation is working correctly. Here is a jupyter notebook with a visual of what is discussed below elongation_test.zip. There I confirm the results we compute match the analytic result from sympy.

The reason we are getting a wrong value for a_major/a_minor in test_elongation on symmetric grids is a resolution issue. The two grids

fg = LinearGrid(L=2, M=eq.M_grid, N=2, sym=False, NFP=eq.NFP)
sg = LinearGrid(L=2, M=eq.M_grid, N=2, sym=True, NFP=eq.NFP)

place nodes at different locations in $\theta$. To compute the elongation, we compute a perimeter of a constant $\zeta$ surface by integrating $\Vert \mathbf{e}_{\theta} \Vert$ over $\theta$. For the particular surface used in test_elongation, the quadrature estimates the true integral better when we sample the integrand at the nodes used by fg than it does when we sample at the nodes used by sg. At reasonable resolution, this difference is not measurable because our midpoint scheme integration will converge to the true integral exponentially on the periodic domain over $\theta$. Increasing eq.M_grid from 2 to 4 will pass the test.

As a sanity check , after the fixes in #451 , we can construct a symmetric grid that has the same nodes as fg like this:

unique_theta = fg.nodes[fg.unique_theta_idx, 1]
sg_manual = LinearGrid(
    L=2, theta=unique_theta[unique_theta <= np.pi], N=2, sym=True, NFP=eq.NFP
)

I have checked computations on sg_manual match fg at all resolutions.

(some details just for your information: The reason sg does not have the same node placement at sg_manual is just a design choice:
When a number is passed in for L, M, N (or rho, theta, zeta) on LinearGrid, we want to construct a grid where all the nodes are weighted equally in any integration etc. The simplest way to do this is to use node placement such that grid.spacing and grid.weights can be the the same across every node (like it is for fg). When symmetry is set to true, it is only correct for grid.spacing and grid.weights to be the same across every node if the nodes are evenly spaced. This requires that the arc distance between the closest node to $\theta = \pi$ and its reflection across the $\theta = \pi$ axis (let's call this its left neighbor assuming counterclock wise is left) equals the distance to the node to its right. In particular, this means there can't be a node at $\theta = 0$ since that distance is 0, so sg must have a different node placement than fg in order for sg.spacing and sg.weights to be constant.

When the node placement is explicitly specified, as in sg_manual, it seems important to respect that placement and not change anything. For the grid sg_manual, the attributes sg.spacing and sg.weights are not constant - in particular, the nodes at $\theta \neq 0$ have double the weight and spacing as the node at $\theta = 0$ because we need to double count the ones at $\theta \neq 0$ since there would have been effectively been two of them in an integration if we had not deleted the nodes with $\theta > \pi$. This ensures that all computations on sg_manual match those of fg.
)

Let me know if there is another potential issue besides test_elongation. As far as I'm aware, the only bug with grids in DESC is if the user wants to use grids with duplicate nodes to compute integrals. (We ensure that the correct computations are performed in the specific case where endpoint=True fur surface integrals, but not for weird custom node patterns with many duplicates). For example we prioritize correctness of surface integrals over line integrals as mentioned in the docstring here. This could be fixed by creating an alias attribute grid.line_spacing that returns grid.spacing if there are no duplicate nodes or a different array of shape (grid.num_nodes, 3) when there are duplicates.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -0.23 +/- 1.24     | -2.86e-05 +/- 1.55e-04 |  1.25e-02 +/- 1.2e-04  |  1.25e-02 +/- 9.3e-05  |
 test_build_transform_fft_midres         |     +0.87 +/- 1.12     | +8.03e-04 +/- 1.04e-03 |  9.35e-02 +/- 8.8e-04  |  9.27e-02 +/- 5.5e-04  |
 test_build_transform_fft_highres        |     -0.28 +/- 0.72     | -1.32e-03 +/- 3.39e-03 |  4.69e-01 +/- 2.4e-03  |  4.70e-01 +/- 2.4e-03  |
 test_equilibrium_init_lowres            |     +1.30 +/- 1.81     | +1.03e-02 +/- 1.44e-02 |  8.04e-01 +/- 1.0e-02  |  7.93e-01 +/- 1.0e-02  |
 test_equilibrium_init_medres            |     -0.87 +/- 1.26     | -1.23e-02 +/- 1.80e-02 |  1.41e+00 +/- 1.2e-02  |  1.42e+00 +/- 1.3e-02  |
 test_equilibrium_init_highres           |     -0.61 +/- 0.80     | -2.58e-02 +/- 3.37e-02 |  4.18e+00 +/- 2.2e-02  |  4.20e+00 +/- 2.6e-02  |
 test_objective_compile_dshape_current   |     +0.89 +/- 8.32     | +3.87e-02 +/- 3.61e-01 |  4.38e+00 +/- 2.9e-01  |  4.34e+00 +/- 2.2e-01  |
 test_objective_compile_atf              |     -0.64 +/- 6.53     | -5.90e-02 +/- 6.04e-01 |  9.18e+00 +/- 3.3e-01  |  9.24e+00 +/- 5.0e-01  |
 test_objective_compute_dshape_current   |     +0.92 +/- 2.81     | +1.96e-05 +/- 5.97e-05 |  2.15e-03 +/- 4.1e-05  |  2.13e-03 +/- 4.3e-05  |
 test_objective_compute_atf              |     +2.57 +/- 1.66     | +1.90e-04 +/- 1.23e-04 |  7.59e-03 +/- 1.0e-04  |  7.40e-03 +/- 7.1e-05  |
 test_objective_jac_dshape_current       |     -0.28 +/- 8.99     | -1.32e-04 +/- 4.27e-03 |  4.73e-02 +/- 3.1e-03  |  4.75e-02 +/- 2.9e-03  |
 test_objective_jac_atf                  |     +1.62 +/- 7.92     | +3.47e-02 +/- 1.70e-01 |  2.18e+00 +/- 1.4e-01  |  2.15e+00 +/- 1.0e-01  |
 test_perturb_1                          |     -3.53 +/- 13.52    | -3.13e-01 +/- 1.20e+00 |  8.54e+00 +/- 8.7e-01  |  8.86e+00 +/- 8.2e-01  |
 test_perturb_2                          |     +0.63 +/- 5.83     | +9.16e-02 +/- 8.51e-01 |  1.47e+01 +/- 7.1e-01  |  1.46e+01 +/- 4.7e-01  |

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (6e8c29f) 95.15% compared to head (3da8572) 95.14%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #807      +/-   ##
==========================================
- Coverage   95.15%   95.14%   -0.01%     
==========================================
  Files          80       80              
  Lines       19371    19378       +7     
==========================================
+ Hits        18432    18438       +6     
- Misses        939      940       +1     

Files	Coverage Δ
desc/compute/_geometry.py	81.57% <ø> (ø)
desc/equilibrium/equilibrium.py	96.31% <100.00%> (+0.03%)	⬆️

... and 1 file with indirect coverage changes

[dpanici]

I have a question on how master_compute_data.pkl was updated, was the only failing key there the elongation and that is why it was updated? did you check that the other keys were still the same?

[dpanici]

can you add to the test_elongation test function a case where you pass in the override grid, so that that part of the compute logic gets tested? Looks fine otherwise

[ddudt]

I have a question on how master_compute_data.pkl was updated, was the only failing key there the elongation and that is why it was updated? did you check that the other keys were still the same?

Yes the only failing key was elongation, since the dimension was changed from a scalar to a vector.

[ddudt]

can you add to the test_elongation test function a case where you pass in the override grid, so that that part of the compute logic gets tested? Looks fine otherwise

This is already covered by the plot_1d test of elongation. The patch coverage for this PR is 100% according to codecov