added setup scripts

root_gnn_dgl/README.md

@@ -0,0 +1,41 @@
+# root_gnn_dgl
+
+## Environment Setup
+
+The environment dependencies for this project are listed in `environment.yml`. Follow the steps below to set up the environment:
+
+### Step 1: Install Conda
+If you don’t already have Conda installed, install either Miniconda (lightweight) or Anaconda (full version):
+
+- **Miniconda**: Download and install from [https://docs.conda.io/en/latest/miniconda.html](https://docs.conda.io/en/latest/miniconda.html).
+- **Anaconda**: Download and install from [https://www.anaconda.com/products/distribution](https://www.anaconda.com/products/distribution).
+
+### Step 2: Clone the Repository
+Clone this repository to your local machine:
+```bash
+git lfs install
+git clone https://huggingface.co/HWresearch/GNN4Colliders
+```
+If you want to clone without large files - just their pointers
+```bash
+GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/HWresearch/GNN4Colliders
+```
+
+### Step 3: Create the Conda Environment
+Use the `environment.yml` file to create the Conda environment:
+```bash
+conda env create -f setup/environment.yml -n <environment_name>
+```
+
+### Step 4: Activate the Environment
+Activate the newly created environment:
+```bash
+conda activate <environment_name>
+```
+Replace <environment_name> with the name of the environment specified in Step 4.
+
+### Step 5: Test the Environment
+Run the `setup/test_setup.py` script to confirm that all packages needed for training are properly set up.
+```bash
+python setup/test_setup.py
+```

root_gnn_dgl/scripts/training_script.py

@@ -29,8 +29,6 @@ import torch.multiprocessing as mp
 from torch.utils.data.distributed import DistributedSampler
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-print("import time: {:.4f} s".format(time.time() - start_time))
-
 def mem():
     print(f'Current memory usage: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 / 1024} GB')
 

root_gnn_dgl/setup/environment.yml

@@ -0,0 +1,391 @@
+name: dgl
+channels:
+  - pytorch
+  - dglteam/label/cu118
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1
+  - _openmp_mutex=4.5
+  - _sysroot_linux-64_curr_repodata_hack=3
+  - afterimage=1.21
+  - anyio=3.7.1
+  - appdirs=1.4.4
+  - argon2-cffi=21.3.0
+  - argon2-cffi-bindings=21.2.0
+  - arrow=1.2.3
+  - asttokens=2.2.1
+  - async-lru=2.0.4
+  - atk-1.0=2.38.0
+  - attrs=23.1.0
+  - awkward-pandas=2023.8.0
+  - aws-c-auth=0.7.0
+  - aws-c-cal=0.6.0
+  - aws-c-common=0.8.23
+  - aws-c-compression=0.2.17
+  - aws-c-event-stream=0.3.1
+  - aws-c-http=0.7.11
+  - aws-c-io=0.13.28
+  - aws-c-mqtt=0.8.14
+  - aws-c-s3=0.3.13
+  - aws-c-sdkutils=0.1.11
+  - aws-checksums=0.1.16
+  - aws-crt-cpp=0.20.3
+  - aws-sdk-cpp=1.10.57
+  - babel=2.12.1
+  - backcall=0.2.0
+  - backports=1.0
+  - backports.functools_lru_cache=1.6.5
+  - beautifulsoup4=4.12.2
+  - binutils=2.38
+  - binutils_impl_linux-64=2.38
+  - binutils_linux-64=2.38.0
+  - blas=1.0
+  - bleach=6.0.0
+  - brotlipy=0.7.0
+  - bzip2=1.0.8
+  - c-ares=1.19.1
+  - c-compiler=1.5.2
+  - ca-certificates=2025.4.26
+  - cached-property=1.5.2
+  - cached_property=1.5.2
+  - cairo=1.16.0
+  - certifi=2024.8.30
+  - cffi=1.15.1
+  - cfitsio=4.2.0
+  - charset-normalizer=2.0.4
+  - comm=0.1.4
+  - compilers=1.5.2
+  - cryptography=41.0.2
+  - cuda-cudart=11.8.89
+  - cuda-cupti=11.8.87
+  - cuda-libraries=11.8.0
+  - cuda-nvrtc=11.8.89
+  - cuda-nvtx=11.8.86
+  - cuda-runtime=11.8.0
+  - cxx-compiler=1.5.2
+  - davix=0.8.4
+  - debugpy=1.6.8
+  - decorator=5.1.1
+  - defusedxml=0.7.1
+  - dgl=1.1.1.cu118
+  - entrypoints=0.4
+  - exceptiongroup=1.1.3
+  - executing=1.2.0
+  - expat=2.5.0
+  - ffmpeg=4.3
+  - fftw=3.3.10
+  - filelock=3.9.0
+  - flit-core=3.9.0
+  - font-ttf-dejavu-sans-mono=2.37
+  - font-ttf-inconsolata=3.000
+  - font-ttf-source-code-pro=2.038
+  - font-ttf-ubuntu=0.83
+  - fontconfig=2.14.2
+  - fonts-conda-ecosystem=1
+  - fonts-conda-forge=1
+  - fortran-compiler=1.5.2
+  - fqdn=1.5.1
+  - freetype=2.12.1
+  - fribidi=1.0.10
+  - ftgl=2.4.0
+  - gcc=11.2.0
+  - gcc_impl_linux-64=11.2.0
+  - gcc_linux-64=11.2.0
+  - gdk-pixbuf=2.42.8
+  - gettext=0.21.1
+  - gflags=2.2.2
+  - gfortran=11.2.0
+  - gfortran_impl_linux-64=11.2.0
+  - gfortran_linux-64=11.2.0
+  - giflib=5.2.1
+  - gl2ps=1.4.2
+  - glew=2.1.0
+  - glog=0.6.0
+  - gmp=6.2.1
+  - gmpy2=2.1.2
+  - gnutls=3.6.15
+  - graphite2=1.3.13
+  - graphviz=6.0.2
+  - gsl=2.7
+  - gsoap=2.8.123
+  - gtk2=2.24.33
+  - gts=0.7.6
+  - gxx=11.2.0
+  - gxx_impl_linux-64=11.2.0
+  - gxx_linux-64=11.2.0
+  - harfbuzz=7.3.0
+  - icu=72.1
+  - idna=3.4
+  - importlib-metadata=6.8.0
+  - importlib-resources=6.0.1
+  - importlib_metadata=6.8.0
+  - importlib_resources=6.0.1
+  - intel-openmp=2023.1.0
+  - ipykernel=6.25.1
+  - ipyparallel=8.6.1
+  - ipython=8.12.2
+  - isoduration=20.11.0
+  - jedi=0.19.0
+  - jinja2=3.1.2
+  - jpeg=9e
+  - json5=0.9.14
+  - jsonpointer=2.0
+  - jsonschema=4.19.0
+  - jsonschema-specifications=2023.7.1
+  - jsonschema-with-format-nongpl=4.19.0
+  - jupyter-lsp=2.2.0
+  - jupyter_client=8.3.0
+  - jupyter_core=5.3.0
+  - jupyter_events=0.7.0
+  - jupyter_server=2.7.0
+  - jupyter_server_terminals=0.4.4
+  - jupyterlab=4.0.5
+  - jupyterlab_pygments=0.2.2
+  - jupyterlab_server=2.24.0
+  - kernel-headers_linux-64=3.10.0
+  - keyutils=1.6.1
+  - krb5=1.20.1
+  - lame=3.100
+  - lcms2=2.12
+  - ld_impl_linux-64=2.38
+  - lerc=3.0
+  - libabseil=20230125.3
+  - libarrow=12.0.1
+  - libblas=3.9.0
+  - libbrotlicommon=1.0.9
+  - libbrotlidec=1.0.9
+  - libbrotlienc=1.0.9
+  - libcblas=3.9.0
+  - libcrc32c=1.1.2
+  - libcublas=11.11.3.6
+  - libcufft=10.9.0.58
+  - libcufile=1.7.1.12
+  - libcurand=10.3.3.129
+  - libcurl=8.1.2
+  - libcusolver=11.4.1.48
+  - libcusparse=11.7.5.86
+  - libcxx=15.0.7
+  - libcxxabi=15.0.7
+  - libdeflate=1.12
+  - libedit=3.1.20191231
+  - libev=4.33
+  - libevent=2.1.12
+  - libexpat=2.5.0
+  - libffi=3.4.4
+  - libgcc-devel_linux-64=11.2.0
+  - libgcc-ng=13.1.0
+  - libgd=2.3.3
+  - libgfortran-ng=11.2.0
+  - libgfortran5=11.2.0
+  - libglib=2.76.4
+  - libglu=9.0.0
+  - libgomp=13.1.0
+  - libgoogle-cloud=2.12.0
+  - libgrpc=1.56.2
+  - libiconv=1.17
+  - libidn2=2.3.4
+  - libllvm13=13.0.1
+  - libllvm14=14.0.6
+  - libnghttp2=1.52.0
+  - libnpp=11.8.0.86
+  - libnsl=2.0.0
+  - libnuma=2.0.18
+  - libnvjpeg=11.9.0.86
+  - libpng=1.6.39
+  - libprotobuf=4.23.3
+  - librsvg=2.54.4
+  - libsodium=1.0.18
+  - libsqlite=3.42.0
+  - libssh2=1.11.0
+  - libstdcxx-devel_linux-64=11.2.0
+  - libstdcxx-ng=13.1.0
+  - libtasn1=4.19.0
+  - libthrift=0.18.1
+  - libtiff=4.4.0
+  - libtool=2.4.7
+  - libunistring=0.9.10
+  - libutf8proc=2.8.0
+  - libuuid=2.38.1
+  - libwebp=1.2.4
+  - libwebp-base=1.2.4
+  - libxcb=1.15
+  - libxml2=2.10.4
+  - libzlib=1.2.13
+  - llvmlite=0.40.1
+  - lz4-c=1.9.4
+  - markupsafe=2.1.1
+  - matplotlib-inline=0.1.6
+  - metakernel=0.29.5
+  - mistune=3.0.0
+  - mkl=2023.1.0
+  - mkl-service=2.4.0
+  - mkl_fft=1.3.6
+  - mkl_random=1.2.2
+  - mpc=1.1.0
+  - mpfr=4.0.2
+  - mpmath=1.3.0
+  - nbclient=0.8.0
+  - nbconvert-core=7.7.3
+  - nbformat=5.9.2
+  - ncurses=6.4
+  - nest-asyncio=1.5.6
+  - nettle=3.7.3
+  - networkx=3.1
+  - nlohmann_json=3.11.2
+  - notebook=7.0.2
+  - notebook-shim=0.2.3
+  - numba=0.57.1
+  - numpy=1.24.3
+  - numpy-base=1.24.3
+  - openh264=2.1.1
+  - openssl=3.3.1
+  - orc=1.9.0
+  - overrides=7.4.0
+  - packaging=23.0
+  - pandas=2.0.3
+  - pandocfilters=1.5.0
+  - pango=1.50.14
+  - parso=0.8.3
+  - pcre=8.45
+  - pcre2=10.40
+  - pexpect=4.8.0
+  - pickleshare=0.7.5
+  - pillow=9.4.0
+  - pip=23.2.1
+  - pixman=0.40.0
+  - pkgutil-resolve-name=1.3.10
+  - platformdirs=2.6.0
+  - pooch=1.4.0
+  - portalocker=2.7.0
+  - prometheus_client=0.17.1
+  - prompt-toolkit=3.0.39
+  - prompt_toolkit=3.0.39
+  - psutil=5.9.0
+  - pthread-stubs=0.4
+  - ptyprocess=0.7.0
+  - pure_eval=0.2.2
+  - pyarrow=12.0.1
+  - pycparser=2.21
+  - pygments=2.16.1
+  - pyopenssl=23.2.0
+  - pysocks=1.7.1
+  - pythia8=8.309
+  - python=3.8.17
+  - python-dateutil=2.8.2
+  - python-fastjsonschema=2.18.0
+  - python-json-logger=2.0.7
+  - python-tzdata=2024.2
+  - python_abi=3.8
+  - pytorch=2.0.1
+  - pytorch-cuda=11.8
+  - pytorch-mutex=1.0
+  - pytz=2023.3
+  - pyyaml=6.0
+  - pyzmq=25.1.1
+  - rdma-core=28.9
+  - re2=2023.03.02
+  - readline=8.2
+  - referencing=0.30.2
+  - requests=2.31.0
+  - rfc3339-validator=0.1.4
+  - rfc3986-validator=0.1.1
+  - root=6.28.0
+  - root_base=6.28.0
+  - rpds-py=0.9.2
+  - s2n=1.3.46
+  - scipy=1.10.1
+  - scitokens-cpp=0.7.3
+  - send2trash=1.8.2
+  - setuptools=68.0.0
+  - six=1.16.0
+  - snappy=1.1.10
+  - sniffio=1.3.0
+  - soupsieve=2.3.2.post1
+  - sqlite=3.41.2
+  - stack_data=0.6.2
+  - sympy=1.11.1
+  - sysroot_linux-64=2.17
+  - tbb=2021.8.0
+  - terminado=0.17.1
+  - tinycss2=1.2.1
+  - tk=8.6.12
+  - tomli=2.0.1
+  - torchaudio=2.0.2
+  - torchtriton=2.0.0
+  - torchvision=0.15.2
+  - tornado=6.3.2
+  - tqdm=4.65.0
+  - traitlets=5.9.0
+  - typing_extensions=4.12.2
+  - typing_utils=0.1.0
+  - ucx=1.14.1
+  - uri-template=1.3.0
+  - urllib3=1.26.16
+  - vdt=0.4.3
+  - vector-classes=1.4.3
+  - wcwidth=0.2.6
+  - webcolors=1.13
+  - webencodings=0.5.1
+  - websocket-client=1.6.1
+  - wheel=0.38.4
+  - xorg-fixesproto=5.0
+  - xorg-kbproto=1.0.7
+  - xorg-libice=1.1.1
+  - xorg-libsm=1.2.4
+  - xorg-libx11=1.8.6
+  - xorg-libxau=1.0.11
+  - xorg-libxcursor=1.2.0
+  - xorg-libxdmcp=1.1.3
+  - xorg-libxext=1.3.4
+  - xorg-libxfixes=5.0.3
+  - xorg-libxft=2.3.8
+  - xorg-libxpm=3.5.16
+  - xorg-libxrender=0.9.11
+  - xorg-libxt=1.3.0
+  - xorg-renderproto=0.11.1
+  - xorg-xextproto=7.3.0
+  - xorg-xproto=7.0.31
+  - xrootd=5.5.4
+  - xxhash=0.8.1
+  - xz=5.2.6
+  - yaml=0.2.5
+  - zeromq=4.3.4
+  - zipp=3.16.2
+  - zlib=1.2.13
+  - zstd=1.5.2
+  - pip:
+      - awkward==2.6.4
+      - awkward-cpp==33
+      - contourpy==1.1.0
+      - cramjam==2.8.3
+      - cycler==0.11.0
+      - fonttools==4.42.0
+      - fsspec==2024.3.1
+      - h5py==3.9.0
+      - pip-install==1.3.5
+      - joblib==1.3.2
+      - kiwisolver==1.4.4
+      - matplotlib==3.7.2
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==8.9.2.26
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - nvidia-cusolver-cu12==11.4.5.107
+      - nvidia-cusparse-cu12==12.1.0.106
+      - nvidia-nccl-cu12==2.20.5
+      - nvidia-nvjitlink-cu12==12.4.127
+      - nvidia-nvtx-cu12==12.1.105
+      - pyparsing==3.0.9
+      - scikit-learn==1.3.0
+      - threadpoolctl==3.2.0
+      - torch==2.3.0
+      - triton==2.3.0
+      - typing-extensions==4.11.0
+      - tzdata==2024.1
+      - uproot==5.3.7
+prefix: /global/homes/j/joshuaho/.conda/envs/dgl

root_gnn_dgl/setup/test_setup.py

@@ -0,0 +1,48 @@
+import os
+import importlib.util
+import sys
+
+def test_imports(directories):
+    """
+    Test importing all Python files in the specified directories.
+    
+    Parameters:
+    - directories: List of directory paths to test.
+    """
+    print("Testing Conda environment...")
+
+    for directory in directories:
+        print(f"\nChecking directory: {directory}")
+        
+        # Check if the directory exists
+        if not os.path.isdir(directory):
+            print(f"Directory not found: {directory}")
+            continue
+        
+        # Iterate through all files in the directory
+        for filename in os.listdir(directory):
+            # Only consider Python files
+            if filename.endswith(".py"):
+                filepath = os.path.join(directory, filename)
+                module_name = os.path.splitext(filename)[0]  # Remove .py extension
+                
+                try:
+                    # Dynamically import the module
+                    spec = importlib.util.spec_from_file_location(module_name, filepath)
+                    module = importlib.util.module_from_spec(spec)
+                    spec.loader.exec_module(module)
+                    print(f"Successfully imported: {filepath}")
+                except Exception as e:
+                    # Print the file and the error message if import fails
+                    print(f"Failed to import: {filepath}")
+                    print(f"Error: {e}")
+
+if __name__ == "__main__":
+    # Automatically append the current directory to sys.path
+    current_directory = os.getcwd()
+    sys.path.append(current_directory)
+    print(f"Current directory added to sys.path: {current_directory}")
+
+    # List of directories to check
+    directories = ["scripts", "root_gnn_base", "models"]
+    test_imports(directories)