En/review eg/dynamics fix

[eduardolneto]

Adding sharding safe check for number of devices, adding capability of initializing particles at a given distance from an axis, correcting energy/pitch/mu/perpnedicular_velocity methods for different types of tracing and adding first iteration of differentiable loss functions for loss fraction, heat flux and loss positions. Changes the dynamics.py.

[rogeriojorge]

Mostly looks good! A few things I spot:

1. dynamics.py:1371-1375, loss_fraction_collisions_differentiable(): this checks hasattr(self, "energy") and then uses self.energy as if it were an array, but energy is a method. This should likely be energies = self.energy() once, then use energies * per_particle_loss[:, None].

2. dynamics.py:1414-1427, escape_location_rmax(): this uses trajectories = self.trajectories, so the weighted escape locations can include the full state dimension, not just position. Since the docstring and penalty functions expect (n_timesteps, 3), this should use trajectories_xyz = self.trajectories[:, :, :3] for the positions.

3. dynamics.py:1512, 1585, 1659, and 1741, classifier penalty methods: the JIT static args look wrong/incomplete. For example, escape_location_penalty_classifier() marks target_position static but not boundary. The target/line/plane/band arrays should stay dynamic, while boundary should be static, e.g. @partial(jit, static_argnums=(0,), static_argnames=["boundary"]). Otherwise this can fail with non-array boundary objects or recompile unnecessarily for target arrays.

4. dynamics.py:1335-1336and1391-1392: the detailed differentiable loss functions compute cumsum(...)and then normalize by the final value, which makes the final loss approximately1` whenever there is any loss. Is this intended? I think these should return the actual cumulative loss fraction/probability (divide by total number of particles?)

Minor things

1. remove the commented-out sharding block at dynamics.py:21-27, fix the indentation at 31-35, remove unused xm_axis and random_phis_offset in InitializeParticlesAroundSurfaceAxis() (209, 236), and avoid reusing the original PRNG key for initial_vparallel_over_v at 288 after splitting it.

2. Another minor thing: dynamics.py:1128-1129, v_perp() full-orbit branch: jnp.sqrt(vperp_squared) can produce NaNs from tiny negative roundoff. Please use jnp.sqrt(jnp.maximum(vperp_squared, 0.0)).

3. Final minor thing. dynamics.py:1268-1269, loss_fraction_collisions(): this currently indexes with x-1, so the first particle uses index -1, and lost_indices[x-1]-1 can also wrap to the last timestep. I’d replace this with the other wrapper thing already used in loss_fraction_BioSavart_collisions()

[eduardolneto]

Addressed all issues here. I removed the differentiable loss functions as they will not be used in the present examples, and they are better added at a future pull request for alpha collisional transport optimzation.

[rogeriojorge]

Looks good to merge

[EstevaoMGomes]

Overall, it looks good. I think the PR is ready to be merged.