switch to tfds array_record, improve visualization, dataset descripti…

…ons (#190)

* improve visualization

* habana vis

* update plots

* fix horovod training

* update click vis

* update visualization

* swap back y and z

---------

Co-authored-by: Javier Duarte <jduarte@ucsd.edu>
Former-commit-id: 1176df0

.github/workflows/test.yml

@@ -10,12 +10,12 @@ on:
 
 jobs:
   deps:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
   deps-pyg:
@@ -25,110 +25,110 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
       - run: pip3 install torch==1.13.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
       - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric==2.2.0 -f https://data.pyg.org/whl/torch-1.13.0+cpu.html
 
   tf-clic-pipeline:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     needs: [deps]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          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-20.04
+    runs-on: ubuntu-22.04
     needs: [deps]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          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-20.04
+    runs-on: ubuntu-22.04
     needs: [deps]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          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-20.04
+    runs-on: ubuntu-22.04
     needs: [deps]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
       - run: ./scripts/local_test_cms_pipeline.sh
 
   pyg-cms-pipeline:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     needs: [deps-pyg]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
       - run: pip3 install torch==1.13.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
       - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric==2.2.0 -f https://data.pyg.org/whl/torch-1.13.0+cpu.html
       - run: ./scripts/local_test_pyg_cms.sh
 
   pyg-delphes-pipeline:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     needs: [deps-pyg]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
       - run: pip3 install torch==1.13.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
       - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric==2.2.0 -f https://data.pyg.org/whl/torch-1.13.0+cpu.html
       - run: ./scripts/local_test_pyg_delphes.sh
 
   pyg-clic-pipeline:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     needs: [deps-pyg]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
       - run: pip3 install torch==1.13.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
       - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric==2.2.0 -f https://data.pyg.org/whl/torch-1.13.0+cpu.html
       - run: ./scripts/local_test_pyg_clic.sh
 
   pyg-ssl-pipeline:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     needs: [deps-pyg]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.8.10'
+          python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
       - run: pip3 install torch==1.13.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu

.gitignore

@@ -29,3 +29,9 @@ nohup.out
 slurm-*.out
 
 .vscode
+
+models
+
+*.root
+
+logs

.pre-commit-config.yaml

@@ -46,4 +46,4 @@ repos:
     # 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']
+           '--extend-ignore=E203,E402,W605']

apptainer/python_base.txt

@@ -25,6 +25,8 @@ imageio-ffmpeg
 ipyparallel
 isort
 jupyter
+jupyterlab
+kaleido
 line_profiler
 lmfit
 lz4
@@ -45,6 +47,7 @@ pre-commit
 pyarrow
 pydot
 pygraphviz
+pyhf
 pymultinest
 pynbody
 pytest

apptainer/python_tf.txt

@@ -12,6 +12,5 @@ tensorflow-datasets
 tensorflow-estimator
 tensorflow-model-optimization
 tensorflow-text
-tf-models-official
 tf2onnx
 transformers

apptainer/tf-2.13.0.singularity

@@ -3,11 +3,12 @@ Bootstrap: docker
 From: tensorflow/tensorflow:2.13.0-gpu-jupyter
 
 %files
-    python_base.txt /opt/python_base.txt
-    python_tf.txt /opt/python_tf.txt
+    specs/python_base.txt /opt/python_base.txt
+    specs/python_tf.txt /opt/python_tf.txt
 
 %post
     apt update -y --fix-missing
+    DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt install -y tzdata
     apt install -y make cmake parallel gcc g++ gfortran binutils
     apt install -y libblas3 libblas-dev liblapack3 liblapack-dev libatlas3-base libatlas-base-dev
     apt install -y libtcmalloc-minimal4
@@ -19,6 +20,8 @@ From: tensorflow/tensorflow:2.13.0-gpu-jupyter
     python3 -m pip install -r /opt/python_base.txt
     python3 -m pip install -r /opt/python_tf.txt
     python3 -m pip install hls4ml[profiling]
+    HOROVOD_WITH_TENSORFLOW=1 python3 -m pip install horovod[tensorflow,keras]
+    python3 -m pip install open3d-cpu
 
 %environment
     export PIP_DEFAULT_TIMEOUT=500