Jb/dps - Differential Particle Solver

[joaopedrobiu6]

This adds an Objective ParticleTracer that uses the guiding center equations of motion to trace particles' trajectories given a DESC Equilibrium. The goal of this objective is to optimize the equilibrium directly from quantities measured from the trajectories traced, e.g. the deviation from the initial flux surface.
It currently supports one-particle tracing and optimization of the equilibrium from that trajectory, and multiple-particle tracing. The optimization of the equilibrium from the multi-particle tracing is not yet fixed

TO DO LIST:

• Run with fwd mode for deriv mode, and see if that fixes the NaNs with multiple particles
• Run with the same particle twice, see if NaN occurs (make sure the conditions used worked with a single particle)
• Look at the trajectories for the particles in the case where the optimization gets NaN on the first iteration. Did some exit the equilibrium (psi > 1)?
• How to deal with particle trajectories that exit plasma (psi > 1)? Probably should just cut off integration there, since past the boundary any field info from DESC is nonsense and we do not want to be integrating the particles there
• Make a tutorial notebook showing how to use
• Move from jax.experimental to diffrax

This was developed by @joaopedrobiu6, @dpanici, @f0uriest and @rogeriojorge.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -1.54 +/- 8.94     | -8.55e-03 +/- 4.98e-02 |  5.48e-01 +/- 4.8e-02  |  5.56e-01 +/- 1.3e-02  |
 test_equilibrium_init_medres            |     +0.19 +/- 2.79     | +7.54e-03 +/- 1.10e-01 |  3.96e+00 +/- 8.4e-02  |  3.95e+00 +/- 7.1e-02  |
 test_equilibrium_init_highres           |     -0.68 +/- 2.13     | -3.34e-02 +/- 1.04e-01 |  4.86e+00 +/- 5.9e-02  |  4.90e+00 +/- 8.6e-02  |
 test_objective_compile_dshape_current   |     +1.95 +/- 6.05     | +5.80e-02 +/- 1.80e-01 |  3.03e+00 +/- 1.2e-01  |  2.98e+00 +/- 1.3e-01  |
-test_objective_compute_dshape_current   |    +39.76 +/- 2.64     | +1.44e-03 +/- 9.52e-05 |  5.05e-03 +/- 4.4e-05  |  3.61e-03 +/- 8.5e-05  |
 test_objective_jac_dshape_current       |     +5.12 +/- 7.23     | +1.99e-03 +/- 2.81e-03 |  4.09e-02 +/- 2.3e-03  |  3.89e-02 +/- 1.6e-03  |
-test_perturb_2                          |     +6.30 +/- 1.77     | +9.92e-01 +/- 2.78e-01 |  1.67e+01 +/- 1.1e-01  |  1.57e+01 +/- 2.6e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.43 +/- 1.37     | +7.42e-02 +/- 2.38e-01 |  1.75e+01 +/- 8.1e-02  |  1.74e+01 +/- 2.2e-01  |
 test_proximal_freeb_jac                 |     +0.87 +/- 1.58     | +5.13e-02 +/- 9.32e-02 |  5.94e+00 +/- 6.3e-02  |  5.89e+00 +/- 6.9e-02  |
 test_solve_fixed_iter_compiled          |     +8.58 +/- 3.92     | +1.55e+00 +/- 7.08e-01 |  1.96e+01 +/- 3.5e-01  |  1.81e+01 +/- 6.2e-01  |
-test_LinearConstraintProjection_build   |    +33.62 +/- 3.46     | +2.84e+00 +/- 2.92e-01 |  1.13e+01 +/- 2.1e-01  |  8.44e+00 +/- 2.0e-01  |
 test_objective_compute_ripple_spline    |    +13.39 +/- 7.96     | +4.41e-02 +/- 2.62e-02 |  3.73e-01 +/- 2.5e-02  |  3.29e-01 +/- 6.9e-03  |
 test_objective_grad_ripple_spline       |     +2.18 +/- 3.66     | +3.26e-02 +/- 5.48e-02 |  1.53e+00 +/- 2.8e-02  |  1.50e+00 +/- 4.7e-02  |
 test_build_transform_fft_midres         |     -0.45 +/- 7.85     | -2.50e-03 +/- 4.35e-02 |  5.52e-01 +/- 3.8e-02  |  5.55e-01 +/- 2.1e-02  |
 test_build_transform_fft_highres        |     +2.30 +/- 3.87     | +1.94e-02 +/- 3.26e-02 |  8.62e-01 +/- 1.6e-02  |  8.42e-01 +/- 2.9e-02  |
 test_equilibrium_init_lowres            |     +5.48 +/- 3.07     | +1.97e-01 +/- 1.10e-01 |  3.78e+00 +/- 7.1e-02  |  3.59e+00 +/- 8.4e-02  |
 test_objective_compile_atf              |     +6.04 +/- 2.76     | +3.68e-01 +/- 1.68e-01 |  6.46e+00 +/- 1.1e-01  |  6.10e+00 +/- 1.3e-01  |
-test_objective_compute_atf              |    +67.09 +/- 7.78     | +6.06e-03 +/- 7.02e-04 |  1.51e-02 +/- 6.8e-04  |  9.03e-03 +/- 1.8e-04  |
 test_objective_jac_atf                  |     -0.52 +/- 2.08     | -9.30e-03 +/- 3.72e-02 |  1.78e+00 +/- 1.9e-02  |  1.79e+00 +/- 3.2e-02  |
 test_perturb_1                          |     +1.68 +/- 4.59     | +2.23e-01 +/- 6.11e-01 |  1.35e+01 +/- 3.4e-01  |  1.33e+01 +/- 5.0e-01  |
 test_proximal_jac_atf                   |     -0.25 +/- 1.57     | -1.95e-02 +/- 1.25e-01 |  7.91e+00 +/- 4.0e-02  |  7.93e+00 +/- 1.2e-01  |
 test_proximal_freeb_compute             |     +6.22 +/- 3.17     | +1.03e-02 +/- 5.24e-03 |  1.75e-01 +/- 3.7e-03  |  1.65e-01 +/- 3.7e-03  |
 test_solve_fixed_iter                   |     +0.57 +/- 3.87     | +1.63e-01 +/- 1.11e+00 |  2.89e+01 +/- 8.1e-01  |  2.87e+01 +/- 7.6e-01  |
 test_objective_compute_ripple           |     +0.06 +/- 1.78     | +1.71e-03 +/- 5.06e-02 |  2.84e+00 +/- 3.7e-02  |  2.84e+00 +/- 3.4e-02  |
 test_objective_grad_ripple              |     -2.86 +/- 2.05     | -2.16e-01 +/- 1.55e-01 |  7.34e+00 +/- 3.3e-02  |  7.55e+00 +/- 1.5e-01  |

[dpanici]

Close once #1820 is in

[YigitElma]

Closing as it is completed in #1820