compute quantities along a fieldline

[unalmis]

• Resolves How to compute quantities along a fieldline #719
• Resolves Separate coordinate dependencies from grid dependencies for computable quantities #692
• Resolves Error when computing Poloidal flux and Boozer toroidal stream function #664
• Resolves eq.compute doesn't correctly override grids to compute quantities that require evaluating functions along field lines #904
• Resolves Errors when compute things on custom grids #1077

Needed to compute neoclassical stuff.
Can look at the Gamma_c branch to see how these changes are used.

• According to jax, jnp.take(x, idx, unique_indices=True) is more efficient than x[idx], so this PR makes this change in the grid.compress() method. (The bounce branch needs jnp.take for different reasons, which is the actual motivation for adding jnp.take and its numpy equivalent to desc.backend. )
• Adds an optional resolution requirement string to the register compute function decorator to resolve 664.
• Adds stuff to use eq.compute method for things whose compute functions assume the input data is evaluated along a field line.
• Changes some grid attributes (backwards-compatible) for the grids to work better in different coordinate systems.
• Fixes a bug on LinearGrid with grid.num_zeta = 1 and grid.NFP!=1 if used for a surface integral (surface averages or others such as surface_integral(x) / surface_integral(y) were unaffected since the bug would net a constant factor difference).

[codecov]

Codecov Report

Attention: Patch coverage is 92.05882% with 27 lines in your changes missing coverage. Please review.

Project coverage is 95.23%. Comparing base (9c63c54) to head (f090f5a).
Report is 1882 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/grid.py	91.35%	16 Missing ⚠️
desc/equilibrium/coords.py	28.57%	5 Missing ⚠️
desc/compute/utils.py	97.00%	3 Missing ⚠️
desc/equilibrium/equilibrium.py	91.17%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1024      +/-   ##
==========================================
- Coverage   95.28%   95.23%   -0.06%     
==========================================
  Files          87       87              
  Lines       21772    21908     +136     
==========================================
+ Hits        20745    20863     +118     
- Misses       1027     1045      +18     

Files with missing lines	Coverage Δ
desc/backend.py	89.80% <100.00%> (+0.13%)	⬆️
desc/compute/_bootstrap.py	100.00% <ø> (ø)
desc/compute/_equil.py	100.00% <ø> (ø)
desc/compute/_field.py	100.00% <ø> (ø)
desc/compute/_geometry.py	82.78% <ø> (ø)
desc/compute/_omnigenity.py	100.00% <ø> (ø)
desc/compute/_profiles.py	99.14% <ø> (ø)
desc/compute/_stability.py	100.00% <ø> (ø)
desc/compute/data_index.py	97.05% <100.00%> (+0.39%)	⬆️
desc/utils.py	92.12% <100.00%> (+0.06%)	⬆️
... and 4 more

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +3.43 +/- 8.45     | +1.86e-02 +/- 4.60e-02 |  5.62e-01 +/- 4.1e-02  |  5.44e-01 +/- 2.1e-02  |
 test_build_transform_fft_midres         |     +3.62 +/- 3.52     | +2.25e-02 +/- 2.19e-02 |  6.45e-01 +/- 1.8e-02  |  6.22e-01 +/- 1.3e-02  |
 test_build_transform_fft_highres        |     +1.89 +/- 3.68     | +1.93e-02 +/- 3.76e-02 |  1.04e+00 +/- 3.2e-02  |  1.02e+00 +/- 1.9e-02  |
 test_equilibrium_init_lowres            |     +5.41 +/- 3.29     | +2.12e-01 +/- 1.29e-01 |  4.13e+00 +/- 8.1e-02  |  3.92e+00 +/- 1.0e-01  |
 test_equilibrium_init_medres            |     +4.06 +/- 5.02     | +1.80e-01 +/- 2.22e-01 |  4.61e+00 +/- 7.9e-02  |  4.43e+00 +/- 2.1e-01  |
 test_equilibrium_init_highres           |     +3.66 +/- 2.73     | +2.12e-01 +/- 1.58e-01 |  6.00e+00 +/- 9.4e-02  |  5.79e+00 +/- 1.3e-01  |
 test_objective_compile_dshape_current   |     +2.43 +/- 1.78     | +9.59e-02 +/- 7.03e-02 |  4.05e+00 +/- 5.0e-02  |  3.95e+00 +/- 4.9e-02  |
 test_objective_compile_atf              |     +2.50 +/- 3.37     | +2.10e-01 +/- 2.84e-01 |  8.62e+00 +/- 2.6e-01  |  8.41e+00 +/- 1.1e-01  |
 test_objective_compute_dshape_current   |     +2.61 +/- 7.36     | +3.32e-05 +/- 9.36e-05 |  1.30e-03 +/- 8.4e-05  |  1.27e-03 +/- 4.2e-05  |
 test_objective_compute_atf              |     +4.91 +/- 10.44    | +2.23e-04 +/- 4.73e-04 |  4.76e-03 +/- 3.8e-04  |  4.53e-03 +/- 2.8e-04  |
 test_objective_jac_dshape_current       |     -4.08 +/- 9.03     | -1.57e-03 +/- 3.47e-03 |  3.68e-02 +/- 2.2e-03  |  3.84e-02 +/- 2.6e-03  |
 test_objective_jac_atf                  |     +5.20 +/- 2.59     | +9.59e-02 +/- 4.78e-02 |  1.94e+00 +/- 4.3e-02  |  1.84e+00 +/- 2.0e-02  |
 test_perturb_1                          |     +3.34 +/- 1.89     | +4.55e-01 +/- 2.58e-01 |  1.41e+01 +/- 1.4e-01  |  1.36e+01 +/- 2.2e-01  |
 test_perturb_2                          |     +2.29 +/- 1.75     | +4.29e-01 +/- 3.27e-01 |  1.92e+01 +/- 2.0e-01  |  1.87e+01 +/- 2.6e-01  |
 test_proximal_jac_atf                   |     +1.16 +/- 1.61     | +8.48e-02 +/- 1.17e-01 |  7.39e+00 +/- 7.2e-02  |  7.30e+00 +/- 9.2e-02  |
 test_proximal_freeb_compute             |     +0.01 +/- 1.39     | +2.61e-05 +/- 2.52e-03 |  1.81e-01 +/- 2.3e-03  |  1.81e-01 +/- 1.1e-03  |
 test_proximal_freeb_jac                 |     +2.39 +/- 1.57     | +1.78e-01 +/- 1.17e-01 |  7.62e+00 +/- 1.1e-01  |  7.44e+00 +/- 4.8e-02  |
 test_solve_fixed_iter                   |     +0.16 +/- 4.25     | +2.91e-02 +/- 7.89e-01 |  1.86e+01 +/- 5.3e-01  |  1.85e+01 +/- 5.9e-01  |

[unalmis]

Approving but wait till after the next release to merge so we can try to keep things organized a bit.

That's fine, but fyi if you release a new version without this PR, we will get complaints from #1077. Recall the current pip version does not have #985 .

[f0uriest]

One question about this and related PRs:

Do we plan to add logic to eq.compute for "default" grids for field line quantities? IE, if a user just asks for epsilon_effective at rho=[0.1, 0.2, ...] but doesn't specify alpha/zeta, can we do that? Or do we require them to pass in a full grid in rho, alpha, zeta coordinates?

[unalmis]

The latter. Feel free to create a new github issue if that'd be convenient

I think I've incroporated Dario's suggestions, and the new DESC version has already been released so I'm planning to merge now.