Summary

ML-based particle flow (MLPF) focuses on developing full event reconstruction for particle detectors using computationally scalable and flexible machine learning models. The project aims to improve particle flow reconstruction across various detector environments, including CMS, as well as future detectors via Key4HEP. We build on existing, open-source simulation software by the experimental collaborations.

Read the user documentation for the physics overview, installation and workflow guides, current capabilities, publications, and roadmap. The Markdown source is in docs/.

---

TLDR; I just want to run the code

You can use uv to set up the repo and test that everything works:

git clone --recurse-submodules https://github.com/jpata/particleflow.git
uv sync --project envs/ort-gpu
uv run ./scripts/local_test_cld.sh
uv run ./scripts/local_test_cms.sh

Alternatively, you can use a prepared container:

apptainer exec --nv https://jpata.web.cern.ch/jpata/pytorch-20260305-08d6950.sif ./scripts/local_test_cld.sh
apptainer exec --nv https://jpata.web.cern.ch/jpata/pytorch-20260305-08d6950.sif ./scripts/local_test_cms.sh

Datasets

See the user documentation to choose a dataset, download and verify published TFDS data, or produce and publish a dataset.

Training

The training guide provides a verified CPU check, a bounded GPU example, expected outputs, and checkpoint continuation.

Model Upload

To upload a trained model to the Hugging Face Hub:

uv run python3 scripts/upload_model_hf.py experiments/pyg-clic-hits-v1_clic_20260328_144021_479374 --version v3.1.0

Model Download & Evaluation

The Key4HEP evaluation guide downloads a published checkpoint and runs it on an example EDM4hep ROOT file. The validation overview explains which checks are needed before interpreting physics performance.

End-to-end workflow

The dataset-generation guide documents the Pixi/Snakemake pipeline from detector simulation through validated TFDS output. Training, evaluation, physics validation, and ONNX validation are covered in the user documentation.

---

Publications

The following publications trace the development of MLPF from early proofs of concept to full detector simulations and fine-tuning studies across detectors.

• [2021] First full-event GNN demonstration of MLPF: Paper Code Dataset
• [2021] First demonstration in CMS Run 3: Paper CMS-DP
• [2022] Improved performance in CMS Run 3: CMS-DP
• [2024] Improved performance with full simulation for future colliders: Paper Code Results
• [2025] Fine-tuning across detectors: Paper Code
• [2026] CMS Run 3 full results: Paper CMS-DP Code

---

Citations and Reuse

You are welcome to reuse the code in accordance with the LICENSE.

How to Cite

1. Academic Work: Please cite the specific papers listed in the Publications section above relevant to the method you are using (e.g., initial GNN idea, fine-tuning, or specific detector studies).
2. Code Usage: If you use the code significantly for research, please cite the specific tagged version from Zenodo.
3. Dataset Usage: Cite the appropriate dataset via the Zenodo link and the corresponding paper.

Contact

For other collaboration ideas, please get in touch via GitHub Discussions.