Skip to content

Machine Learning for Particle Track Reconstruction

Welcome to repository and documentation for ML pipelines and techniques by the ExatrkX Collaboration. Here we present a set of templates, best practices and results gathered from significant trial and error, to speed up the development of others in the domain of machine learning for high energy physics. We focus on applications specific to detector physics, but many tools can be applied to other areas, and these are collected in an application-agnostic way in the Tools section.


To start as quickly as possible, clone the repository, Install and follow the steps in Quickstart. This will get you generating toy tracking data and running inference immediately. Many of the choices of structure will be made clear there. If you already have a particle physics problem in mind, you can apply the Template that is most suitable to your use case.

Once up and running, you may want to consider more complex ML Models. Many of these are built on other libraries (for example Pytorch Geometric).


It's recommended to start a conda environment before installation:

conda create --name exatrkx-tracking python=3.8
conda activate exatrkx-tracking
pip install pip --upgrade

If you have a CUDA GPU available, load the toolkit or install it now.


will attempt to negotiate a path through the packages required, using nvcc --version to automatically find the correct wheels. Warning: If you are installing with cpu, this may take up to 15 minutes due to an unfortunately slow installation of Pytorch3D from source.

You should be ready for the Quickstart!

If this doesn't work, you can step through the process manually:

1. Run `export CUDA=cpu` 1a. Find the GPU version cuda XX.X with `nvcc --version` 1b. Run `export CUDA=cuXXX`, with `XXX = 92, 101, 102, 110`
2. Install Pytorch and dependencies ```pip install --user -r requirements.txt```
3. Install local packages ```pip install -e .```
4. Install CPU-optimized packages ```pip install faiss-cpu pip install "git+"``` 4. Install GPU-optimized packages ```pip install faiss-gpu cupy-cudaXXX```, with `XXX` ```pip install pytorch3d -f{Y}_cu{XXX}_pyt{ZZZ}/download.html``` where `{Y}` is the minor version of Python 3.{Y}, `{XXX}` is as above, and `{ZZZ}` is the version of Pytorch {Z.ZZ}. e.g. `py36_cu101_pyt170` is Python 3.6, Cuda 10.1, Pytorch 1.70.