Graph Neural Networks for Charged Particle Tracking

Description

Charged particle tracking reconstructs the trajectories (“tracks”) of elementary particles traveling through a detector. This task is different from many other problems that involve trajectories:

there are several thousand particles that need to be tracked at once,
there is no time information (the particles travel too fast),
we do not observe a continuous trajectory but instead only around five points (“hits”) along the way in different detector layers.

The task can be described as a combinatorically very challenging “connect-the-dots” problem, essentially turning a cloud of points (hits) in 3D space into a set of O(1000) trajectories. Expressed differently, each hit (containing not much more than the x/y/z coordinate) must be assigned to the particle/track it belongs to.

Tracking algorithms are a crucial part of the software used at big high energy physics experiments and have a strong impact on their physics performance. They are also computationally very expensive and can limit the rate at which data taking occurs. For these reasons, tremendous development efforts have been spent on tracking software. Historically, tracking software was mostly specific to individual physics experiments, but recent efforts such as Acts have shown the potential of strong open-source community efforts.

Unlike traditional tracking algorithms that are built around Kalman filters, this project uses graph neural networks for significant increases in speed. A conceptually simple way to turn tracking into a machine learning task is to create a fully connected graph of all points and then train an edge classifier to reject any edge that doesn’t connect points that belong to the same particle. In this way, only the individual trajectories remain as components of the initial fully connected graph. However, this strategy does not lead to perfect results in practice, so it is usually combined with other steps.

While many of the model architectures are still actively explored and a variety of pipelines are under investigation, this project aims to build up a modular code base for many different GNN tracking studies. By providing different models and common utility functions, the framework comes batteries-included while still remaining flexible for new ideas of any kind. Our hope is that the code base will lower the entry-barrier for new studies and become a focal point of the community.

The project code together with documentation and a reading list is available on github.

Objectives and Tasks

Depending on the candidate, two different areas can be explored:

Area 1: Software engineering. Example tasks:

Increase training efficiency and GPU utilization
Port tracking metrics from numpy/pandas to pytorch
Switch to a GPU-ready implementation of DBSCAN
Improve use of hyperparameter optimization scripts (currently built around ray tune) on HPC clusters with the SLURM workload manager
General work on unit testing and maintenance tasks

Area 2: Machine learning & statistics. Example tasks:

Investigate and contribute new model architectures
Implement and investigate different clustering algorithms
Run side-studies for variations of the existing architecture
Run hyperparameter optimizations

Of course, students can also mix tasks from both areas. Many of the concrete tasks are collected as issues in the main GNN Tracking repository.

Requirements

For area 1:

Excellent python skills, ideally including pytorch or numpy
Working knowledge on git
Dream candidate would specifically have experience with pytorch, GPU programming, or ray

For area 2:

High degree of mathematical understanding: candidate should be able to read general machine learning papers and implement the described architectures or loss functions
General experience in data science and machine learning frameworks, preferably with pytorch
Dream candidate would have some prior experience with graph neural networks, ideally with pytorch geometric

Expected results

Example outcomes for area 1:

Significantly faster training, making the best use of the available resources
The whole pipeline should run on the GPU
A user-friendly workflow that runs hyperparameter sweeps with ray tune) on a system with the SLURM workload manager.

Example outcomes for area 2:

Increased overall efficiency/performance of the default pipeline
Detailed writeup of different architectures/algorithms and corresponding results
Code available for every side-study, ideally in the main repository as a model that can be switched out with the current one, alternatively as a verbose Jupyter notebook.

Mentors

Kilian Lieret kl5675@princeton.edu Princeton University
Gage deZoort jdezoort@princeton.edu Princeton University

Please see the instructions below for the phase 1 application.

Phase 1 Application

Phase 1 has concluded. Unfortunately, we do not take any applications anymore.

Please send us your CV and motivation for this project.

We also have a qualification task to make a fair selection of candidates. It should take you no more than 1-2 days to complete (depending on prior knowledge).

For fairness, you will have 48h to send us your solutions after you receive the instructions.

We understand that your schedule might be busy, so we proceed as follows:

Send us a start date and time in Eastern Time zone (EST). Please let us know 1-2 working days in advance to ensure we can schedule the email in time. Your submission should be completed by Mar 12, so the latest start date is Mar 10.
We will send a quick confirmation email within 1-2 working days.
The email with the qualification task instructions will be sent at the time specified in step 1.
You work for 1-2 days and submit your results (submissions after 1 or 2 days will be evaluated separately, so that the speed of completion is taken into account).
We will confirm receiving your solutions within 1-2 working days.
We will announce the result of this phase 1 selection on Mar 14.