Update benchmarking branch with lates developments from main.

.github/workflows/pre-commit.yml

@@ -0,0 +1,27 @@
+name: Run pre-commit
+
+on:
+  pull_request:
+    branches: [ main ]
+  push:
+    branches: [ main ]
+
+jobs:
+  lint:
+    name: Lint PR or Push to main
+    runs-on: ubuntu-20.04
+    strategy:
+      matrix:
+        python-version: [3.8.10]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v2
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Run Lint
+        uses: pre-commit/action@v2.0.0

.github/workflows/test.yml

@@ -2,70 +2,162 @@ name: CI
 
 on:
   push:
-    branches: [ master ]
+    branches: [ main ]
   pull_request:
-    branches: [ master ]
+    branches: [ main ]
 
   workflow_dispatch:
 
 jobs:
-  delphes-pipeline:
-    runs-on: ubuntu-latest
+  deps:
+    runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v2
-      - name: Install python deps
-        run: |
-          sudo apt install python3 python3-pip wget
-          sudo python3 -m pip install --upgrade pip
-          sudo python3 -m pip install --upgrade setuptools
-          sudo python3 -m pip install tensorflow==2.6 setGPU \
-            sklearn matplotlib mplhep pandas scipy uproot3 uproot3-methods \
-            awkward0 keras-tuner networkx \
-            tensorflow-probability==0.12.2 tensorflow-addons==0.13.0 \
-            tqdm click tensorflow-datasets 'ray[default]' 'ray[tune]' \
-            tf-models-official tensorflow-text \
-            tf2onnx onnxruntime zenodo_get seaborn
-          git submodule init
-          git submodule update
-      - name: Run delphes TF model
-        run: ./scripts/local_test_delphes_pipeline.sh
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+  deps-pyg:
+    runs-on: ubuntu-20.04
+    needs: [deps]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: pip3 install torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cpu.html
+
+  tf-unittest:
+    runs-on: ubuntu-22.04
+    needs: [deps]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: ./scripts/local_test_cms_pipeline.sh
+      - run: PYTHONPATH=. python3 -m unittest tests/test_tf.py
+
+  tf-clic-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [tf-unittest]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: ./scripts/local_test_clic_pipeline.sh
+
+  tf-clic-hits-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [tf-unittest]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: ./scripts/local_test_clic_hits_pipeline.sh
+
+  tf-delphes-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [tf-unittest]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: ./scripts/local_test_delphes_pipeline.sh
+
+  tf-cms-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [tf-unittest]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: ./scripts/local_test_cms_pipeline.sh
+
+  pyg-unittests:
+    runs-on: ubuntu-22.04
+    needs: [deps-pyg]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: pip3 install torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cpu.html
+      - run: PYTHONPATH=. python3 -m unittest tests/test_torch_and_tf.py
+
+  pyg-cms-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [pyg-unittests]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: pip3 install torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cpu.html
+      - run: ./scripts/local_test_pyg_cms.sh
+
+  pyg-delphes-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [pyg-unittests]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: pip3 install torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cpu.html
+      - run: ./scripts/local_test_pyg_delphes.sh
 
-  cms-pipeline:
-    runs-on: ubuntu-latest
+  pyg-clic-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [pyg-unittests]
     steps:
-      - uses: actions/checkout@v2
-      - name: Install python deps
-        run: |
-          sudo apt install python3 python3-pip wget
-          sudo python3 -m pip install --upgrade pip
-          sudo python3 -m pip install --upgrade setuptools
-          sudo python3 -m pip install tensorflow==2.6 setGPU \
-            sklearn matplotlib mplhep pandas scipy uproot3 uproot3-methods \
-            awkward0 keras-tuner networkx \
-            tensorflow-probability==0.12.2 tensorflow-addons==0.13.0 \
-            tqdm click tensorflow-datasets 'ray[default]' 'ray[tune]' \
-            tf-models-official tensorflow-text \
-            tf2onnx onnxruntime zenodo_get seaborn
-          git submodule init
-          git submodule update
-      - name: Run CMS TF model using the pipeline
-        run: ./scripts/local_test_cms_pipeline.sh
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: pip3 install torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cpu.html
+      - run: ./scripts/local_test_pyg_clic.sh
 
-  delphes-pytorch:
-    runs-on: ubuntu-latest
+  pyg-ssl-pipeline:
+    runs-on: ubuntu-22.04
+    needs: [pyg-unittests]
     steps:
-      - uses: actions/checkout@v2
-      - name: Install python deps
-        run: |
-          sudo apt install python3 python3-pip wget
-          sudo python3 -m pip install --upgrade pip
-          sudo python3 -m pip install --upgrade setuptools
-          sudo python3 -m pip install setGPU sklearn matplotlib mplhep pandas scipy uproot3 uproot3-methods awkward0 networkx
-          sudo python3 -m pip install torch==1.6.0 torchvision==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
-          sudo python3 -m pip install --no-index torch-scatter -f https://pytorch-geometric.com/whl/torch-1.6.0+cpu.html
-          sudo python3 -m pip install --no-index torch-sparse -f https://pytorch-geometric.com/whl/torch-1.6.0+cpu.html
-          sudo python3 -m pip install --no-index torch-cluster -f https://pytorch-geometric.com/whl/torch-1.6.0+cpu.html
-          sudo python3 -m pip install --no-index torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.6.0+cpu.html
-          sudo python3 -m pip install torch-geometric comet_ml
-      - name: Run Delphes Pytorch model
-        run: ./scripts/local_test_delphes_pytorch.sh
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+      - run: pip install -r requirements.txt
+      - run: pip3 install torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cpu.html
+      - run: ./scripts/local_test_pyg_ssl.sh

.gitignore

@@ -3,22 +3,34 @@
 *.npz
 *.pt
 *.pdf
+*.png
 data/*
+experiments/*
+prp/*
 *.pth
-test/__pycache__/
 
-mlpf/pytorch/__pycache__/*
-mlpf/plotting/__pycache__/*
-mlpf/pytorch/data
-test_tmp/
-test_tmp_delphes/
+*/__pycache__/*
 .DS_Store
 
-prp
 *.pyc
 *.pyo
 
-mlpf/updated/LRP/pid*
-mlpf/updated/LRP/class*
-
 *.ipynb_checkpoints
+
+*playground.py
+*playground.ipynb
+nohup.out
+
+*.pkl
+*.pkl.bz2
+*.parquet
+
+slurm-*.out
+
+.vscode
+
+models
+
+*.root
+
+logs

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "hep_tfds"]
-    path = hep_tfds
-    url = https://github.com/jpata/hep_tfds

.pre-commit-config.yaml

@@ -0,0 +1,49 @@
+default_language_version:
+    python: python3
+exclude: ^(scripts/delphes/tev14_pythia8_*.)|^(images)|^(scripts/fcc/clicRec_e4h_input.py)|^(scripts/fcc/clic_steer.py)|^(scripts/fcc/pythia.py)|^(scripts/fcc/PandoraSettings/.*)
+
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.4.0
+  hooks:
+  - id: trailing-whitespace
+  - id: check-added-large-files
+  - id: check-ast
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-xml
+  - id: check-yaml
+  - id: debug-statements
+  - id: end-of-file-fixer
+  - id: requirements-txt-fixer
+  - id: mixed-line-ending
+    args: ['--fix=no']
+
+#This is giving different results on github CI vs. locally and I can't understand why
+#- repo: https://github.com/PyCQA/isort
+#  rev: 5.11.2
+#  hooks:
+#  - id: isort
+#    args: ['--profile', 'black', '--filter-files']
+
+- repo: https://github.com/psf/black
+  rev: 22.12.0
+  hooks:
+    - id: black
+      # It is recommended to specify the latest version of Python
+      # supported by your project here, or alternatively use
+      # pre-commit's default_language_version, see
+      # https://pre-commit.com/#top_level-default_language_version
+      language_version: python3
+      args: [--line-length=125]
+
+- repo: https://github.com/PyCQA/flake8
+  rev: 6.0.0
+  hooks:
+  - id: flake8
+    # black-compatible flake-8 config
+
+    # E203 is not PEP8 compliant
+    # E402 due to logging.basicConfig in pipeline.py
+    args: ['--max-line-length=125',  # github viewer width
+           '--extend-ignore=E203,E402,W605']

README.md

@@ -1,28 +1,36 @@
-[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.4559587.svg)](https://doi.org/10.5281/zenodo.4559587)
-[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.4559324.svg)](https://doi.org/10.5281/zenodo.4559324)
 [![CI](https://github.com/jpata/particleflow/workflows/CI/badge.svg)](https://github.com/jpata/particleflow/actions)
 
+# Overview
+MLPF focuses on developing full event reconstruction based on computationally scalable and flexible end-to-end ML models.
+
 <p float="left">
-  <img src="notebooks/plots/event.png" alt="Simulated event" width="600"/>
+  <img src="images/schematic.png" alt="High-level overview" width="600"/>
 </p>
 
+## MLPF development in CMS
+
 <p float="left">
-  <img src="notebooks/plots/num_particles.png" alt="Particle multiplicity" width="300"/>
-  <img src="notebooks/plots/res_pid2.png" alt="Neutral hadron resolution" width="300"/>
+  <img src="images/cms/ev_pf.png" alt="PF reconstruction" width="300"/>
+  <img src="images/cms/ev_mlpf.png" alt="MLPF reconstruction" width="300"/>
 </p>
 
-## MLPF with Delphes
-
-Short instructions with a single test file in [notebooks/delphes-tf-mlpf-quickstart.ipynb](notebooks/delphes-tf-mlpf-quickstart.ipynb).
+<p float="left">
+  <img src="images/cms/ak4jet_puppi_pt_ttbar.png" alt="PUPPI jets in ttbar" width="300"/>
+</p>
 
-Long instructions for reproducing the full training from scratch in [README_delphes.md](README_delphes.md).
-The plots can be generated using the notebook [delphes/delphes_model_analysis.ipynb](delphes/delphes_model_analysis.ipynb).
+  - ACAT 2022:
+    - CERN-CMS-DP-2022-061, http://cds.cern.ch/record/2842375
+  - ACAT 2021:
+    - J. Phys. Conf. Ser. 2438 012100, http://dx.doi.org/10.1088/1742-6596/2438/1/012100
+    - CERN-CMS-DP-2021-030, https://cds.cern.ch/record/2792320
 
-### Delphes dataset
-The dataset is available from zenodo: https://doi.org/10.5281/zenodo.4559324.
+## Initial development with Delphes
 
-### Software setup
-The software setup for the ML training is available in the singularity spec file [scripts/base.singularity](scripts/base.singularity).
+<p float="left">
+  <img src="images/delphes/num_particles.png" alt="Number of reconstructed particles" width="250"/>
+  <img src="images/delphes/inference_time.png" alt="Scaling of the inference time" width="300"/>
+</p>
 
-## Acknowledgements
-This project is supported by the Mobilitas Pluss Returning Researcher Grant MOBTP187 of the Estonian Research Council. Part of this work was conducted at **iBanks**, the AI GPU cluster at Caltech. We acknowledge NVIDIA, SuperMicro and the Kavli Foundation for their support of **iBanks**. 
+  - paper: https://doi.org/10.1140/epjc/s10052-021-09158-w
+  - code: https://doi.org/10.5281/zenodo.4559587
+  - dataset: https://doi.org/10.5281/zenodo.4559324