Handle VMEC axis events in particle tracing

[rogeriojorge]

What changed

• Add a magnetic-axis termination event at s=1e-6 for the two collisionless VMEC guiding-center models: GuidingCenter and GuidingCenterAdaptative.
• Preserve the existing s=1 loss event and retain its event mask separately from axis terminations.
• Forward-fill only axis-terminated output tails, replacing Diffrax post-event inf saves with the last finite sampled state for diagnostics and plotting.
• Report the number of axis-terminated particles as unresolved.
• Keep the existing loss-fraction calculation, but stop JIT-compiling the complete Tracing instance as a static argument.

Scope

This PR intentionally does not change solver limits, progress-meter defaults, rejected-step handling, or the public constructor API. Those generic controls belong in a separate PR.

The axis event is not applied to full-orbit models: VMEC guiding-center states use flux coordinates (s, theta, phi, v_parallel), while full-orbit states use Cartesian (x, y, z, vx, vy, vz). Treating the first full-orbit coordinate as s would be physically incorrect. Collision and field-line models are also unchanged in this focused fix.

Root cause

VMEC (s, theta, phi) coordinates are singular at s=0. A QI test orbit crossed the magnetic axis, but ESSOS previously terminated only at s=1, so it extrapolated the VMEC representation to negative flux and the solve became increasingly stiff. Separately, loss_fraction treated the complete Tracing instance as a static JIT argument, causing a long post-trace compilation for larger trajectory arrays.

Validation

• python -m pytest tests/test_dynamics.py -q: 15 passed
• QI, 50 particles, tmax=5e-4, 10 CPU devices:
  • GuidingCenter: 6.27 s, 0% lost, 1 axis-unresolved
  • GuidingCenterAdaptative: 12.98 s, 0% lost, 1 axis-unresolved
• The restored trace_particles.py keeps its original structure and completes without stalling.

The full repository suite previously showed 81 passes and 13 unrelated pre-existing failures under JAX 0.11.