Welcome to JetNet!

For developing and reproducing ML + HEP projects.

JetNet • Installation • Quickstart • Documentation • Contributing • Citation • References

CI Documentation Status Codestyle pre-commit.ci status PyPI Version PyPI Downloads DOI

JetNet

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!

Installation

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).

Quickstart

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/", download=True)
# or as a PyTorch dataset:
dataset = TopTagging(jet_type="all", data_dir="./datasets/toptagging/", split="train", download=True)

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()

Documentation

The full API reference and tutorials are available at jetnet.readthedocs.io. Tutorial notebooks are in the tutorials folder, with more to come.

Contributing

We welcome feedback and contributions! Please feel free to create an issue for bugs or functionality requests, or open pull requests from your forked repo to solve them.

Building and testing locally

Perform an editable installation of the package from inside your forked repo and install the pytest package for unit testing:

pip install -e .
pip install pytest

Run the test suite to ensure everything is working as expected:

pytest tests                    # tests all datasets
pytest tests -m "not slow"      # tests only on the JetNet dataset for convenience

Citation

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).

References

[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].