For developing and reproducing ML + HEP projects.
JetNet • Installation • Quickstart • Documentation • Citation • References
JetNet is an effort to increase accessibility and reproducibility in jet-based machine learning.
Currently we provide:
- Easy-to-access and standardised interfaces for the following datasets:
- Standard implementations of generative evaluation metrics (Ref. [1, 2]), including:
- Fréchet physics distance (FPD)
- Kernel physics distance (KPD)
- Wasserstein-1 (W1)
- Fréchet ParticleNet Distance (FPND)
- coverage and minimum matching distance (MMD)
- Loss functions:
- Differentiable implementation of the energy mover's distance [3]
- And more general jet utilities.
Additional functionality is under development, and please reach out if you're interested in contributing!
JetNet can be installed with pip:
pip install jetnet
To use the differentiable EMD loss jetnet.losses.EMDLoss
, additional libraries must be installed via
pip install "jetnet[emdloss]"
Finally, PyTorch Geometric must be installed independently for the Fréchet ParticleNet Distance metric jetnet.evaluation.fpnd
(Installation instructions).
Datasets can be downloaded and accessed quickly, for example:
from jetnet.datasets import JetNet, TopTagging
# as numpy arrays:
particle_data, jet_data = JetNet.getData(jet_type=["g", "q"], data_dir="./datasets/jetnet/")
# or as a PyTorch dataset:
dataset = TopTagging(jet_type="all", , data_dir="./datasets/toptagging/", split="train")
Evaluation metrics can be used as such:
generated_jets = np.random.rand(50000, 30, 3)
fpnd_score = jetnet.evaluation.fpnd(generated_jets, jet_type='g')
Loss functions can be initialized and used similarly to standard PyTorch in-built losses such as MSE:
emd_loss = jetnet.losses.EMDLoss(num_particles=30)
loss = emd_loss(real_jets, generated_jets)
loss.backward()
The full API reference and tutorials are available at jetnet.readthedocs.io. Tutorial notebooks are in the tutorials folder, with more to come.
More detailed information about each dataset can (or will) be found at jet-net.github.io.
If you find this library useful for your research, please consider citing our original paper which introduces it [1].
Additionally, if you use our EMD loss implementation, please cite the respective qpth or cvxpy libraries, depending on the method used (qpth
by default).
[1] R. Kansal et al., Particle Cloud Generation with Message Passing Generative Adversarial Networks, NeurIPS 2021 [2106.11535].
[2] R. Kansal et al., Evaluating Generative Models in High Energy Physics, Phys. Rev. D 107 (2023) 076017 [2211.10295].
[3] P. T. Komiske, E. M. Metodiev, and J. Thaler, The Metric Space of Collider Events, Phys. Rev. Lett. 123 (2019) 041801 [1902.02346].