Graph Neural Networks for Track Finding in Mu3e

Daniel King > Under the Supervision of Nikolaos Rompotis

The Mu3e experiment at the Paul Scherrer Institute searches for the charged-lepton-flavour-violating decay μ⁺ → e⁺ e⁺ e⁻ (signal) with the goal of excluding branching fractions above 10⁻¹⁶ at the 90% confidence level. This decay is essentially forbidden by the Standard Model (SM) and any observation would provide evidence of new physics. The experiment faces a loss in the number of observed MC signal decays (i.e. a lowered efficiency) when applying cuts that purify signal detection. An alternative approach involves a Graph Neural Network (GNN) and finds signal and beam efficiencies similar to the official standard reconstruction algorithm. More information on the experiment is available in the Technical design of the phase I Mu3e experiment by the Mu3e collaboration.

View and interact with Monte Carlo tracks and reconstructed graphs here: https://dan-k-k.github.io/plotly-mu3e/. Here is a screenshot of the true tracks left by a MC μ⁺ → e⁺ e⁺ e⁻ decay.

This repository contains the full implementation of a Graph Neural Network (GNN) model for the Mu3e experiment:

• Implements a six‐hit graph-building procedure with loose geometric and kinematic constraints, pruning away 99.9% combinatorial (fake) graphs.
• Defines and trains a three-class GNN (electron vs. positron vs. fake) using PyTorch Geometric.

The original Mu3e Monte Carlo “.root” files used to build the graphs are not included in this repository.

Report overview:

Monte Carlo simulations accurately model the expected decay particles' hits left in the detector, made up of frames: a period of time long enough to capture all hits from the decay. The Graph Neural Network approach contained in this repository focuses on six-hit tracks that reach the recurl pixel layers. A three-class GNN uses spatial and kinematic information of combinations of hits left in the detector to classify each six-hit graph as an electron, positron, or fake (combinatorial background). Common background decay types are internal conversion μ⁺ → e⁺ e⁺ e⁻ ν_e ν̄_μ and Michel μ⁺ → e⁺ ν_e ν̄_μ.

Graphs are formed by combining hits that follow a detection layer sequence, assuming no track information is known. A cheap constraints algorithm for graph generation greatly reduces processing time, essentially focusing on graphs that closest resemble true particle tracks.

The GNN model is trained and validated on a dataset of ~92,000 six-hit graphs of one signal decay event per frame. In the beam simulation, at most one signal decay is expected to be seen in a frame among the background decays.

The model returns outputs for each class on test graphs (for signal, this is 20% of the available set). The outputs are the model's estimated probabilities of the graph belonging to each class. In some cases, the model will incorrectly give true electron and positron graphs low outputs and vice versa. Note that the y axis is logarithmic.

Electron/positron misclassification is negligible.

There is a trade-off between purity and efficiency. More real and fake graphs are classified real as the required score to pass is lowered. This increases efficiency (more of the true real graphs are correctly recalled), but lowers purity (more fakes are among the real predictions).

For frames with more hits, the model's performance is worse as there is higher chance of nearby hits generating fakes that closely resemble real tracks. For frames with equal hits, the beam set tends to have worse purity due to individual Michel decay events with uniform spatial distributions as opposed to, for example, signal, requiring all electrons to conserve momentum.

Comparing the full GNN method efficiencies for all datasets with the official standard reconstruction and Clopper-Pearson one-sigma errors (with a GNN-approach beam purity of 94.5% and signal purity of 97.7%, while the official standard reconstruction approach is considered pure).

How to run

1. Clone:

   git clone https://github.com/dan-k-k/Mu3e-GNN-code.git
   cd Mu3e-GNN-code

2. Create virtual environment:

   python3 -m venv .venv
   source .venv/bin/activate
   

3. Install requirements:

   pip install -r requirements.txt
   

4. Generate graphs through the Mu3e_Graph_Generation.ipynb notebook

5. Train the model and view results using the Mu3e_GNN_and_Results.ipynb notebook