Permutation Equivariant, Lorentz Invariant/Covariant Aggregator Network for applications in Particle Physics.
At the moment it includes two main variants: a classifier (for e.g. top-tagging) and a 4-momentum regression network.
arXiv link: https://arxiv.org/abs/2211.00454
PELICAN is a network that takes 4-momentum inputs (e.g. jet constituents) and uses Lorentz-invariant features and permutation-equivariant layers to solve tasks such as classification and momentum regression in an exactly Lorentz-covariant way. Embedding the symmetries into the architecture drastically reduces the model size while still delivering state-of-the-art performance. Moreover, by using only physically/geometrically meaningful mathematical operations, PELICAN can provide viable physics models as opposed to black-box ML approaches that do not respect fundamental symmetries. Finally, equivariance massively improves sample efficiency (and generalizability from low amounts of training data). This repository contains two main top-level scripts, one for classification, and one for 4-momentum regression (currently outputting just one 4-vector prediction). These can be executed directly with no installation. The data folder contains a small sample dataset consisting of 50% top quark jets and 50% background jets (with the key is_signal
identifying which is which).
- python >=3.9
- pytorch >=1.10
- h5py
- colorlog
- scikit-learn
- tensorboard (if using --summarize)
- optuna (if using the optuna script)
- psycopg2-binary (if using optuna with a distributed file sharing system)
- No installation required -- just use one of the top level scripts.
- The classifier can be used via the script
train_pelican_classifier.py
. The 4-momentum regressor istrain_pelican_cov.py
. The main required argument is--datadir
since it provides the datasets. - See data/sample_data/ for a small example of datasets. Each datapoint contains some number of input 4-momenta (E,p_x,p_y,p_z) under the key
Pmu
, target 4-momenta (e.g. the true top momentum) undertruth_Pmu
and a classification label underis_signal
. Regression datasets can contain multiple target 4-momenta under keys liketruth_Pmu_0
etc. The choice of the target key is controlled by the argument--target
. - For classification, the number of target classes is set by the argument
--num-classes
. - When num_classes>2, the target is expected to be a one-hot encoding as opposed to a single integer label. Note that the way multi-class metrics are reported is also a bit different.
- The same scripts can be used for inference on the test dataset when it is run with the flag
--task=eval
. - Model checkpoints can be loaded via
--load
for inference or continued training (as long as the prefix is the same). - By default the model is run on a GPU (
--cuda
), but CPU evaluation can be forced via--cpu
. - The argument
--verbose
(correspodinginly--no-verbose
) can be used to write per-minibatch stats into the log. - The number of particles in each event can vary. Most computations are done with masks that properly exclude zero 4-momenta. The argument
--nobj
sets the maximum number of particles loaded from the dataset. Argument--add-beams
also appends two "beam" particles of the form (1,0,0,±1) to the inputs to help the network learn that bias in the dataset due to the fixed z-axis. With this flag, the intermediate tensors in Eq2to2 layers will have shape [Batch, 2+Nobj, 2+Nobj, Channel]. - The network includes learnable rescaling of certain tensors by powers of the ratio
Nobj/Nobj_avg
, whereNobj_avg
is a constant set by--nobj-avg
that should indicate the typical (or average) number of particles in one event. In the large top-tagging dataset used in the paper this number of 49. - The input 4-monenta can be uniformly scaled with a multiplicative
--scale
(default 1.0). The beams are not rescaled. In case of momentum regression, the 4-momentum output by the network is automatically divided by the same scale to better match the scale of the target.
- Logfiles are kept in the
log/
folder. By default these contain initializtion information followed by training and validation stats for each epoch, and testing results at the end. The argument--prefix
is used to name all files. Re-running the script without changing the prefix will overwrite all output files unless the flag--load
is used. - If
--summarize
is on, then Tensorboard summaries are saved into a folder with the same name as the log. If--summarize-scv
is on, then per-minibatch stats are written into a separate CSV file. - At the end of evaluation on the testing set, the stats are written into a CSV file whose name ends in
Best.metrics.csv
(for the model checkpoint with the best validation score) andFinal.metrics.csv
(for the model checkpoint from the last epoch). If there are multiple runs whose prefixes only differ by text after a dash (e.g.run-1, run-2, run-3
, etc.) then their metrics will be appended to the same CSV. - Model outputs (predictions) are saved as .pt files in
predict/
. ROC curves are also saved there as CSV files. In the multiclass case, each class takes three columns of the file (the columns are [fpr, tpr, threshold]). - Model checkpoints are saved as .pt files in
model/
.
- Here is an example of a command that starts training the classifier on the sample dataset that is part of this repository. Optimally there should be three files with names train.h5, valid.h5, and test.h5.
python3 train_pelican_classifier.py --datadir=./data/sample_data/run12 --target=is_signal --nobj=80 --nobj-avg=49 --num-epoch=35 --num-train=60000 --num-valid=60000 --batch-size=64 --prefix=classifier --optim=adamw --activation=leakyrelu --factorize --lr-decay-type=warm --lr-init=0.0025 --lr-final=1e-6 --drop-rate=0.05 --drop-rate-out=0.05 --weight-decay=0.005
- Similarly for regression (this prompt will try to predict the top quark momentum stored in the sample dataset, however since 50% of the dataset are background events without a top quark, for which truth_Pmu=0, this is not a real regression test and will produce subpar predictions):
python3 train_pelican_cov.py --datadir=./data/sample_data/run12 --target=truth_Pmu --nobj=80 --nobj-avg=20 --num-epoch=35 --num-train=60000 --num-valid=60000 --batch-size=64 --prefix=regressor --optim=adamw --activation=leakyrelu --factorize --lr-decay-type=warm --lr-init=0.0025 --lr-final=1e-6 --drop-rate=0.05 --drop-rate-out=0.05 --weight-decay=0.005
Alexander Bogatskiy, Flatiron Institute
Jan T. Offermann, University of Chicago
Timothy Hoffman, University of Chicago
Xiaoyang Liu, University of Chicago
Inspiration, code snippets, etc.
This project is licensed under the MIT License - see the LICENSE.md file for details