Skip to content

mcamp8/heptrkx-gnn-tracking

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

87 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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.

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Jupyter Notebook 99.5%
  • Other 0.5%