Guiding Center Particle Tracing

[YigitElma]

This PR is a combination of #734 and #1410.

The main structure is under discussion in #1657.

• Adds particle tracing capabilities in desc.particles module.
  • Particle tracing is done via desc.particles.trace_particles function.
  • Particles can be initialized in couple different ways:
    • ManualParticleInitializerLab : Initializes particles at given positions in lab coordinates.
    • ManualParticleInitializerFlux : Initializes particles at given positions in flux coordinates.
    • CurveParticleInitializer : Initializes N particles on a given curve.
    • SurfaceParticleInitializer : Initializes N particles on a given surface.
  • Implemented particle trajectory models are:
    • VacuumGuidingCenterTrajectory : Integrates the particle motion by vacuum guiding center ODEs, conserving energy and mu.
  • Particle trajectories can be plotted with desc.plotting.plot_particle_trajectories function.

We can delete #734 and #1410 once this is merged.

[YigitElma]

I was playing around with different PID coefficients, with scaled precise_QA with 20 particles. Looks like pcoeff=0.3, icoeff=0.3, dcoeff=0 takes less number of steps for harder particles, if the integration happens very quickly then it performs a bit worse but those are very fast so doesn't matter.

rtz_1, _, aux_1 = _trace_particles(
    field=eq,
    y0=x0,
    model=model,
    model_args=args,
    ts=ts,
    params=eq.params_dict,
    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, pcoeff=0.4, icoeff=0.3, dcoeff=0
    ),
    saveat=SaveAt(ts=ts),
    solver=Tsit5(),
    adjoint=RecursiveCheckpointAdjoint(),
    event=event,
    chunk_size=None,
    options={},
    throw=False,
    return_aux=True,
)

[f0uriest]

a few minor points, otherwise looks fine

[f0uriest]

Note that this might break in the future, I think @ddudt and @maya-avida were planning on making Equilibrium subclass from MagneticField as part of #996 #

[f0uriest]

though depending on how efficient the implementation is maybe this would be fine since it would still avoid the rootfinding

[YigitElma]

Based on their API, I can adapt the check later.