Add particle tracker in openPMD-viewer

[openpmd]

Add particle tracker in openPMD-viewer

[jgerity]

I'm not sure exactly how you'd be doing the select filtering on the back-end of that object, but if it does not prove to be too expensive, a suggested improvement would be to allow passing a callable for that argument so that the user can implement whatever selection criterion they'd like (vs choosing from a fixed set of criteria) as if they were using the filter() built-in. E.g.,

pt = ParticleTracker(ts, species='electrons', select=momentumSelector, iteration=2000)

def momentumSelector(particle):
    return particle['uz'] > 10 and particle['uz'] < 100

[RemiLehe]

Thanks for this great suggestion.

Actually for implementing the select I was thinking about using the implementation from ts.get_particle within ParticleTracker.
Currently, ts.get_particle can take a dictionary as argument for select. But it would definitely be a great ideat to allow the user to pass a callable, so as to give more flexibility !
But since this change (which would in the end mainly affect ts.get_particle) is somewhat independent of implementing particle tracking, I'll implement it in a separate pull request.

[soerenjalas]

I like your API approach, especially the idea to have a tracker object containing all information of the selected particles. Do you think having NaNs in the arrays can be a problem? For example do all numpy methods like NaNs? Otherwise, I think this would make a convenient API for a really handy tool!

[RemiLehe]

Yeah, I know the NaNs are not ideal since methods like sum, mean, etc. will return a Nan. However, you can use np.nansum and np.nanmean to ignore NaNs in the array (and compute the sum and mean on the values that are not NaNs). But I agree it is inconvenient.

However, I think it is important that the same particles are returned always at the same position within the array, for different iterations. This is what allows users to reconstruct particle trajectories over several iteration. For instance:

pt = ParticleTracker( ts, species='electrons', select={'uz':[10,1000]}, iteration=2000)

N_iterations = len(ts.iterations)
N_particles = pt.N_selected
X_trajectories = np.empty( ( N_iterations, N_particles ) )
for i in range( N_iterations ):
   x,  = ts.get_particles( ['x'], particle_tracker=pt, iteration=ts.iterations[i] )
   X_trajectories[i, :] = x[:]

trajectory_electron0 = X_trajectories[ :, 0 ]
trajectory_electron1 = X_trajectories[ :, 1 ]

However, if I return the same particle always at the same position in the array, I also need to fill the array with something when particles have exited the simulation box at later iterations, and this is where the NaNs are needed.

Does the above make sense? I am definitely open to other solutions that do not use the NaNs. But I think they have to maintain the property that individual trajectories can still be tracked. If you have any suggestions don't hesitate.

[jgerity]

Perhaps you could allow returning an array of ids in the resulting tuple? This lets the user track by ID by keeping track of order in the two separate arrays, and shouldn't cause any breaking change if I've understood the get_particle method right. E.g.:

pid, x, y, z = ts.get_particle( ['id', 'x', 'y', 'z'], particle_tracker=pt, iteration=50000 )

[soerenjalas]

Sorry for bringing the issue up without having a better solution… I completely see @RemiLehe's point and I also forgot about the NaN resistant numpy methods. So I think for most practical purposes having NaNs should be fine and for edge cases one could still filter the NaN with something like:

trajectory_electron = np.where(trajectory_electron==NaN, 0)

Nevertheless, I think being able to get the particle id, as @jgerity said, is also a useful option. But I guess this will be the case anyway when the pid is a quantity of the species.

And one last thing I assume is already covered by you API suggestion but I just want to be sure. For example when tracking the origin of a particle beam it should be possible to fetch the particles before the selection iteration like:

# Select particles at iteration 2000
pt = ParticleTracker( ts, species='electrons', select={'uz':[10,100]}, iteration=2000 )
x1, y1, z1 = ts.get_particle( ['x', 'y', 'z'], particle_tracker=pt, iteration=1000 )

And now that I have made so many unhelpful comments :), just tell me if you want help programming, e.g. the plotter or the slider widget.

[RemiLehe]

@soerenjalas Yes, the API will support fetching iterations that are before the selection iteration.

Regarding coding : I'm actually almost done, I'll do a PR soon. The most useful would be actually if you can continue checking that the choice of API I made is okay with you, and also if you can review the upcoming PR. :)

[RemiLehe]

Regarding the NaNs: in the end, I think I will add an option preserve_particle_index in the ParticleTracker object:

• If preserve_particle_index is True, then the same particle will always have the same index in the output array, so that we can easily track it as explained above. However, this implies that there will be NaN in the output array.
• If preserve_particle_index is False, then there will be no NaNs (the output array will just have fewer elements because of the absent particles). But then it is not guaranteed, that particles are always at the same position. But then, as @jgerity pointed out, the user can still track particles "by hand", using their id. (In fact, outputting the id is already supported in openPMD-viewer).

So the user can choose preserve_particle_index depending on his/her purpose. preserve_particle_index=True is better for tracking, but preserve_particle_index=False is better for computing average quantities (emittance, etc.)

[RemiLehe]

Also, another change regarding the previous API I described: I think that passing the particle tracker under select would be actually nicer for the user, like so:

# Select particles at iteration 2000
pt = ParticleTracker( ts, species='electrons', select={'uz':[10,100]}, iteration=2000 )

# Fetch them at iteration 30000 and 50000
x1, y1, z1 = ts.get_particle( ['x', 'y', 'z'], select=pt, iteration=30000 )

This means that, in get_particle, select can either be a dictionary (with selection rules) or a ParticleTracker object (and soon a callable)>. This will allow methods like get_emittance, etc. to work directly with the particle tracker, without any change in the code.

[RemiLehe]

@jgerity @soerenjalas Let me know if you have any opposition, or further suggestion, to the above API (the preserve_particle_index and passing the ParticleTracker as select). If not, I will try to do a corresponding PR by the end of today.

[jgerity]

preserve_particle_index looks like a good compromise to me! I would think False would be a good default value, but you know better than I do what the typical user will want for this option. No other suggestions or objections.

[ax3l]

I like the API, it's clean and intuitive.

Regarding your NaN question, you could also go (optional?) for numpy.ma masked arrays.