/heptrkx-gnn-tracking

Primary LanguageJupyter Notebook

Graph Neural Networks for particle track reconstruction

This repository contains the PyTorch implementation of the GNNs for particle track reconstruction from CTD 2018: https://arxiv.org/abs/1810.06111.

Contents

The main python scripts for running:

  • prepare.py: the data preparation script which reads TrackML data files, cleans and reduces the data, and writes hit graphs to the filesystem.
  • train.py: the main training script which is steered by configuration file and loads the data, model, and trainer, and invokes the trainer to train the model.

Other stuff:

  • In the scripts directory are SLURM batch scripts for running the jobs on Cori at NERSC.
  • The GNN model code lives in models/gnn.py.
  • The dataset code for reading the prepared hit graphs lives in datasets/hitgraphs.py.
  • The main trainer code for the GNN segment classifier lives in trainers/gnn.py.