Swap vectorization order in bounce integrals

[unalmis]

After the recent refactoring to the Bounce1D class that resulted from #1214, the API is a little too strict for computations like effective ripple etc. where we vectorize the computation over over some dimensions and loop over others to save memory.

This PR changes the expected shape of the pitch angle input to Bounce1D in #854 from (P, M, L) to (M, L, P). With this change, the two leading axes of all inputs to the methods in that class is (M, L).

These changes are tested and already included in downstream branches. I am making new PR instead of directly committing to the bounce branch for people who have already reviewed the bounce PR.

This is better because

1. Easier usage for end users. (Previously, you'd have to manually add trailing axes to pitch angle array).
2. Makes it much simpler to use with JAX's new batched map.
3. Previously we would loop over the pitch angles to save memory. However, this means some computation is repeated because interpax would interpolate multiple times. By looping over the field lines instead and doing the interpolation for all the pitch angles at once, both _bounce_quadrature and interp_to_argmin are faster. (I'm seeing 30% faster speed with similar memory just from computing effective ripple (no optimization), but I don't plan to do any benchmarking to see whether that is from recent changes like Improve coordinate mapping performance #1154 or Remove jit method of objective, directly compile methods #1043 , or others).

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +1.16 +/- 6.68     | +5.92e-03 +/- 3.41e-02 |  5.16e-01 +/- 3.3e-02  |  5.10e-01 +/- 1.0e-02  |
 test_equilibrium_init_medres            |     -0.43 +/- 0.50     | -1.75e-02 +/- 2.00e-02 |  4.00e+00 +/- 1.7e-02  |  4.02e+00 +/- 1.1e-02  |
 test_equilibrium_init_highres           |     +0.06 +/- 0.81     | +3.19e-03 +/- 4.33e-02 |  5.36e+00 +/- 3.9e-02  |  5.35e+00 +/- 1.8e-02  |
 test_objective_compile_dshape_current   |     +0.36 +/- 1.51     | +1.36e-02 +/- 5.70e-02 |  3.80e+00 +/- 5.5e-02  |  3.78e+00 +/- 1.6e-02  |
 test_objective_compute_dshape_current   |     +1.91 +/- 1.82     | +6.51e-05 +/- 6.19e-05 |  3.47e-03 +/- 4.6e-05  |  3.41e-03 +/- 4.1e-05  |
 test_objective_jac_dshape_current       |     -1.88 +/- 8.20     | -7.46e-04 +/- 3.25e-03 |  3.89e-02 +/- 2.7e-03  |  3.97e-02 +/- 1.8e-03  |
 test_perturb_2                          |     -0.23 +/- 1.15     | -3.94e-02 +/- 1.96e-01 |  1.69e+01 +/- 1.7e-01  |  1.70e+01 +/- 9.8e-02  |
 test_proximal_freeb_jac                 |     +0.45 +/- 0.84     | +3.29e-02 +/- 6.17e-02 |  7.36e+00 +/- 3.8e-02  |  7.33e+00 +/- 4.9e-02  |
 test_solve_fixed_iter                   |     -0.19 +/- 59.54    | -9.09e-03 +/- 2.91e+00 |  4.87e+00 +/- 2.0e+00  |  4.88e+00 +/- 2.1e+00  |

[codecov]

Codecov Report

Attention: Patch coverage is 93.87755% with 3 lines in your changes missing coverage. Please review.

Project coverage is 95.39%. Comparing base (c531a82) to head (08e4257).

Files with missing lines	Patch %	Lines
desc/integrals/bounce_integral.py	85.00%	3 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           bounce    #1242   +/-   ##
=======================================
  Coverage   95.39%   95.39%           
=======================================
  Files          95       95           
  Lines       23147    23192   +45     
=======================================
+ Hits        22081    22125   +44     
- Misses       1066     1067    +1     

Files with missing lines	Coverage Δ
desc/integrals/bounce_utils.py	94.41% <100.00%> (+0.55%)	⬆️
desc/integrals/quad_utils.py	100.00% <100.00%> (ø)
desc/integrals/bounce_integral.py	94.23% <85.00%> (-5.77%)	⬇️

... and 8 files with indirect coverage changes