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.

[YigitElma]

Maybe for these kinds of single node Grid creation, we can use a spline for only the radial part of zernike polynomials. The most expensive part of the transform calculation is the recurrence in the _jacobi which is not very efficient for single-node evaluation. We use vmap over r values but for particle tracing it doesn't perform well. The idea is to compute the zernike radial and its derivatives on a dense 1D grid, pass it to the Transform class, and instead of calling the zernike_radial it will evaluate the spline at the required points. The spline fit can be a callable like interpax.Interpolator1D.

r_dense = jnp.linspace(0.01, 1, 1000)
dense_zernike = zernike_radial(r=r_dense, l=l, m=m)
radial_spline = interpax.Interpolator1D(x=r_dense, f=dense_zernike)
# this should replace zernike call in transform.build
transforms = get_transforms(data_keys, eq, grid, jitable=True, radial_spline=radial_spline)

I haven't tested it yet, but it should potentially solve the problem other options had. FourierChebyshevField is good in practice, but it needs 13 set of higher res spectral coefficients, and each particle needs to access these coefficients, so the GPU memory is bottlenecked (hence 1 particle is faster but as you have more particles, it becomes slower). Interpolator3D should work in practice, but I think it is again bottlenecked by the memory.