Fourier-Chebyshev Fit to B for particle tracing

[YigitElma]

• Adds a new convenience class FourierChebyshevField which has a couple of methods to be compatible with optimization,
  • init : only resolution info
  • build : creates grid and transforms
  • fit : given params and profiles (these change during optimization), fits a new Fourier-Chebyshev series.
  • evaluate : Evaluate the series at requested point.

Example usage will be,

interpolater = FourierChebyshevField(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)
interpolater.build(eq)
data_fit = interpolater.fit(eq.params_dict, {"iota": eq.iota})

rtz, _, aux = _trace_particles(
    field=interpolator,  # normally you pass eq here
    y0=x0,
    model=model,
    model_args=args,
    ts=ts,
    params=None,   # normally you pass eq.params_dict here
    max_steps=int((ts[-1] - ts[0]) / 1e-8),
    min_step_size=1e-8,
    stepsize_controller=PIDController(rtol=1e-4, atol=1e-4, dtmin=1e-8),
    saveat=SaveAt(ts=ts),
    solver=Tsit5(),
    adjoint=RecursiveCheckpointAdjoint(),
    event=event,
    chunk_size=None,
    options={},    # this passes extra stuff to field computation, not needed here
    throw=False,
    return_aux=True,
)

Todo:

• trace_particles should initialize FourierChebyshevField and by default use it. Add a flag to use Equilibrium field
• plot_particle_trajectories etc should also be able to use the interpolated field.
• Look at the spectral convergence (ideally lower res for basis vectors might be good, but current code is simpler and with batching might not be inefficient)
• Consider adding spline version too?

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    3.78 %    |     3.813e+03      |     3.957e+03      |    144.13    |       39.78        |       36.26        |
  test_proximal_jac_w7x_with_eq_update   |   -1.27 %    |     6.663e+03      |     6.579e+03      |    -84.80    |       162.72       |       161.50       |
  test_proximal_freeb_jac                |    0.41 %    |     1.316e+04      |     1.321e+04      |    53.79     |       84.83        |       83.43        |
  test_proximal_freeb_jac_blocked        |    0.41 %    |     7.461e+03      |     7.492e+03      |    30.73     |       72.51        |       73.21        |
  test_proximal_freeb_jac_batched        |   -1.25 %    |     7.538e+03      |     7.444e+03      |    -94.36    |       72.91        |       73.10        |
  test_proximal_jac_ripple               |   -2.37 %    |     3.608e+03      |     3.523e+03      |    -85.50    |       65.70        |       65.71        |
  test_proximal_jac_ripple_bounce1d      |   -0.08 %    |     3.566e+03      |     3.564e+03      |    -2.77     |       76.55        |       76.55        |
  test_eq_solve                          |    0.60 %    |     1.992e+03      |     2.004e+03      |    12.00     |       93.79        |       94.69        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[codecov]

Codecov Report

❌ Patch coverage is 77.69231% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.67%. Comparing base (3753ff0) to head (7b21b40).

Files with missing lines	Patch %	Lines
desc/particles.py	76.98%	29 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2012      +/-   ##
==========================================
- Coverage   95.75%   95.67%   -0.09%     
==========================================
  Files         102      102              
  Lines       28344    28471     +127     
==========================================
+ Hits        27142    27239      +97     
- Misses       1202     1232      +30     

Files with missing lines	Coverage Δ
desc/compute/_basis_vectors.py	100.00% <100.00%> (ø)
desc/particles.py	88.15% <76.98%> (-4.34%)	⬇️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[YigitElma]

For some reason, this suffers from GPU synchronisation miserably, and for more particles, it is slower than direct approach.