MXNet implementation of the jet classification network in ParticleNet: Jet Tagging via Particle Clouds.
The ParticleNet model can be obtained by calling the get_particle_net function in particle_net.py, which can return either an MXNet Symbol or an MXNet Gluon HybridBlock. The model takes three input arrays:
points: the coordinates of the particles in the (eta, phi) space. It should be an array with a shape of (N, 2, P), where N is the batch size and P is the number of particles.features: the features of the particles. It should be an array with a shape of (N, C, P), where N is the batch size, C is the number of features, and P is the number of particles.mask: a mask array with a shape of (N, 1, P), taking a value of 0 for padded positions.
To have a simple implementation for batched training on GPUs, we use fixed-length input arrays for all the inputs, although in principle the ParticleNet architecture can handle variable number of particles in each jet. Zero-padding is used for the points and features inputs such that they always have the same length, and a mask array is used to indicate if a position is occupied by a real particle or by a zero-padded value.
The implementation of a simplified model, ParticleNet-Lite, is also provided and can be accessed with the get_particle_net_lite function.
If you use ParticleNet in your research, please cite the paper:
@article{Qu:2019gqs,
author = "Qu, Huilin and Gouskos, Loukas",
title = "{ParticleNet: Jet Tagging via Particle Clouds}",
year = "2019",
eprint = "1902.08570",
archivePrefix = "arXiv",
primaryClass = "hep-ph",
SLACcitation = "%%CITATION = ARXIV:1902.08570;%%"
}
The ParticleNet model is developed based on the Dynamic Graph CNN model. The implementation of the EdgeConv operation in MXNet is adapted from the author's TensorFlow implementation, and also inspired by the MXNet implementation of PointCNN.