Graph neural networks for classification

As the final step in our array of machine learning techniques for top quark classification we turn towards graph neural networks.

This type of networks is well suited to our problem, as we can describe the event very well using a graph. There are four top quarks (four nodes) which can have correlations (the edges between them). Both nodes and edges can have features. We start with the simplest case possible and consider only the node features pt, eta, phi, and mass for each top quark, as well as the conections to the other three top quarks. Later we can also provide features to the edges, such as the DeltaR variable or perhaps (for trying it out) the invariant mass between the two top quarks denoted by the two connected nodes.

Prerequisites

You need to install additional packages in order to run pytorch-geometric, which we will use to design the graph neural networks.

Please set up the virtual environment (see previous sections if you don't know how to do that).

source setup.sh

You should now have the virtual environment active (visible by the (venv) tag in your command line prompt).

Second, install the missing packages, before moving on with exploring the script for graph-neural-network-based classification.

pip install -r requirements.txt --upgrade

Now you are fully set up to work with pytorch-geometric.

Graph neural network script

A preliminary version of a graph neural network script is provided here:

Beware that it is not working, as besides the dataset being ready as a graph, the network itself is just hastily copied from examples and not really adapted to the problem nor working properly for the moment.

You can try if you get it to run! Feel free to make modifications to the code as you see fit.

Last updated