Batched jet clustering for physics analysis and PyTorch training loops. FlashJet implements anti-kt, kt, and Cambridge/Aachen with NumPy, an optional C++ CPU kernel, PyTorch, and Triton GPU kernels.
Documentation · GitHub releases · Source · Authors
FlashJet is available from Git, with documentation published from main.
Tagged GitHub releases are separate from the current development snapshot.
Install from Git with PyTorch for the batch example below:
python -m pip install 'flashjet[torch] @ git+https://github.com/jet-universe/FlashJet.git'
# Linux x86_64 with a compatible NVIDIA driver and CUDA-enabled PyTorch:
python -m pip install 'flashjet[triton,data] @ git+https://github.com/jet-universe/FlashJet.git'For NumPy-only use, omit [torch]. See the installation guide
for optional dependencies, native builds, and wheels from CI.
Append @<full-commit-sha> to the Git URL to pin an exact revision. PyTorch and
Triton remain optional dependencies. There is no PyPI publication in this setup.
import torch
import flashjet
p4 = torch.tensor([[[10., 0., 0., 10.],
[2., 0., 0., 2.],
[-5., 0., 0., 5.]]], requires_grad=True)
mask = torch.ones(p4.shape[:2], dtype=torch.bool)
out = flashjet.cluster(p4, mask, R=0.4, algorithm="antikt")
jets = out.jets_p4(p4)
jets, order = out.sort_jets_by_pt(jets)
jets[..., 3].sum().backward()The last axis is (px, py, pz, E). Use p4.to("cuda") and a mask on the same
device for GPU clustering. Gradients pass through the sum of constituent
momenta, not through the clustering decisions.
For one event, flashjet.cluster(numpy_array) returns a ClusterSequenceRef.
Call inclusive_jets() for momenta or jet_constituents() for particle indices.
Read the quickstart, plain-English guide, backend limits, and release checklist.
python -m pip install -e '.[torch,test,data]'
pytest -q
python -m pip install -r docs/requirements.txt
sphinx-build -W --keep-going -b html docs docs/_build/htmlCUDA tests skip when no GPU is available. Release automation builds wheels for CPython 3.9–3.13 on Linux x86_64 and macOS Intel/Apple Silicon, plus a source archive. Installation and artifact tests run in CI. See the release checklist for workflow coverage and the separate GPU validation requirements.
GPL-3.0-or-later; see LICENSE. The clustering algorithm derives from FastJet. Cite its papers in physics publications. Authors: Sitian Qian, Chirayu Gupta (@Chirayu18), and Alexandre De Moor (@AlexDeMoor), with AI co-authors OpenAI Codex and Anthropic Claude. See AUTHORS.md.