ParticleTracker mechanics refactoring (again!)

[StanczakDominik]

Tackles, but does not close, #936 just a little.

It's officially Q2 of 2021 and I already feel the roadmap backlog pressure (thank you once again, @rocco8773 😉 ), so I'm starting out with the particle tracker refactoring again. I brought over changes from #675 and made them at least a little more self contained. At this point I have a bad reimplementation of a terrible first implementation.

What I intend to do with this from now on:

• assimilate the changes the wise @pheuer made in diagnostics/proton_radiography.py (some great ideas and an actual use case! that really helps flesh out a code 😆 )
• make this code actually nice
• allow using this with the Grids classes; probably refactor some old cruft out (looking at you, Plasma3D, your names would be numbered if I had an ETA for this)
• make the current proton_radiography use the ParticleTracker class for the actual motion within EM fields part.
  • this means I need to provide a way of running this that, rather than the full/snapshotted trajectory, outputs the final state of these particles. I probably means another method, run_lite or some such, that provides a simpler interface without all the trajectory saving.
  • Allowing variable stopping conditions and incorporating not simulating particles that run out of the grid will thus be necessary.
• Merge the ExB drift notebook with the one I've got for this, since the implementation is actually starting to look nice
• Change the (dumb, in retrospect) idea I had of providing an xarray accessor to a set of viz functions that take xarrays. Should have been obvious.

I also recently had the idea of using google/jax with this. I'm not sure how much of an issue (as usual...) fields calculations would be with that; but I'd like to at least try. Automatically running this on GPUs would be pretty sweet. But that's definitely a future PR.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

Merging #1096 (92cb8a7) into master (5dd2547) will decrease coverage by 0.77%.
The diff coverage is 70.09%.

@@            Coverage Diff             @@
##           master    #1096      +/-   ##
==========================================
- Coverage   96.91%   96.13%   -0.78%     
==========================================
  Files          70       71       +1     
  Lines        6928     7084     +156     
==========================================
+ Hits         6714     6810      +96     
- Misses        214      274      +60     

Impacted Files	Coverage Δ
plasmapy/diagnostics/proton_radiography.py	100.00% <ø> (ø)
plasmapy/simulation/particle_integrators.py	100.00% <ø> (ø)
plasmapy/simulation/particletracker.py	68.91% <67.21%> (-29.37%)	⬇️
plasmapy/plasma/sources/analyticalfields.py	95.00% <95.00%> (ø)
plasmapy/plasma/sources/__init__.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 5dd2547...92cb8a7. Read the comment docs.

[pheuer]

Sounds great, looking forward to this! Looking for any easy optimization opportunities in grids.py is on my to-do list. We can definitely change or add functionality to grids as you need it.

A few comments/ideas on top of your list above:

• I would aim for three output modes: full trajectories, partial trajectories (eg. every N time steps), and only the final positions. It would also be good to allow storing trajectories in some way other than memory (eg. an ascii or hdf5 file) which would be necessary for all but trivial calculations.

• It would be nice for users to be able to provide their own "stop condition" functions (in addition to a few built-in options). I'm thinking the user could pass in a func that takes the particle tracker as its only argument and return a boolean based on whatever internal logic it likes.

• In order to play nicely with grids, we should refactor the particle integrators to take arrays of field scalars rather than vectors (eg. Ex, Ey, Ez instead of E). That would avoid doing this on every push:

 # Create arrays of E and B as required by push algorithm
E = np.array([Ex.to(u.V / u.m).value, Ey.to(u.V / u.m).value, Ez.to(u.V / u.m).value] )
E = np.moveaxis(E, 0, -1)
B = np.array([Bx.to(u.T).value, By.to(u.T).value, Bz.to(u.T).value])
B = np.moveaxis(B, 0, -1)

• Related to that last point, we should create methods for grids that return values without fetching and applying units (just to have the units then stripped off again before going into the particle integrator). For example, the interpolators should have an option to not apply units before returning.

• Don't forget we also have Added relativistic boris push #979 to wrap up at some point!

In addition, here's some longer-term brainstorming.

• Eventually I'd really like to find (or make) 3D interpolators that are faster than the SciPy ones, and that are parallelizable and numba-compatible.

• Another idea I've been thinking about is extending grids to allow a time axis. This would allow, for example, proton radiographs of fields that evolve on the scale of the proton crossing time (capturing motion blurring, for example). I'd guess this would be helpful for particle tracker too.

<sarcasm> I don't think anyone other than cryptocurrency miners are able to get GPUs anymore. Maybe it would be better to optimize around some more ubiquitous architecture. Maybe smart speakers.</sarcasm>

[StanczakDominik]

Just don't look at the code yet! :D

[pheuer]

@StanczakDominik what's the status of this PR in light of #2245?