Merge branch 'graphnet-team:main' into add-MODynEdge

.github/workflows/build.yml

@@ -74,7 +74,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9, '3.10']
+        python-version: [3.8, 3.9, '3.10', '3.11']
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
@@ -88,7 +88,7 @@ jobs:
       - name: Run unit tests and generate coverage report
         run: |
           set -o pipefail  # To propagate exit code from pytest
-          coverage run --source=graphnet -m pytest tests/ --ignore=tests/data/
+          coverage run --source=graphnet -m pytest tests/ --ignore=tests/data/ --ignore=tests/deployment/
           coverage report -m
 
   build-macos:
@@ -108,5 +108,5 @@ jobs:
       - name: Run unit tests and generate coverage report
         run: |
           set -o pipefail  # To propagate exit code from pytest
-          coverage run --source=graphnet -m pytest tests/ --ignore=tests/data/
+          coverage run --source=graphnet -m pytest tests/ --ignore=tests/data/ --ignore=tests/deployment/
           coverage report -m

.gitignore

@@ -149,6 +149,7 @@ data/examples/output/
 **.pth
 **.db
 **.parquet
-
+# Exception to pre-trained folders; here we want all .pth files 
+!/graphnet/src/graphnet/models/pretrained/**/**/**/**/**.pth
 # Notebooks
 **.ipynb

CONTRIBUTING.md

@@ -8,7 +8,7 @@ Use [GitHub issues](https://github.com/graphnet-team/graphnet/issues) for tracki
 
 ## Pull requests
 
-Develop code in a forks of the [main repo](https://github.com/graphnet-team/graphnet). Make contributions in dedicated development/feature branches on your forked repositories, e.g. if you are implementing a specific `GraphBuiler` class you could create a branch named `add-euclidean-graph-builder` on your own fork.
+Develop code in a fork of the [main repo](https://github.com/graphnet-team/graphnet). Make contributions in dedicated development/feature branches on your forked repositories, e.g. if you are implementing a specific `GraphDefinition` class you could create a branch named `add-euclidean-graph-definition` on your own fork.
 
 Create pull requests from your development branch into `graphnet-team/graphnet:main` to contribute to the project. **To be accepted,** pull requests must:
   * pass all automated checks,
@@ -20,7 +20,7 @@ Create pull requests from your development branch into `graphnet-team/graphnet:m
 
 ## Conventions
 
-This repository aims to support python 3 version that are actively supported (currently `>=3.6`). Standard python coding conventions should be followed:
+This repository aims to support python 3 version that are actively supported (currently `>=3.8`). Standard python coding conventions should be followed:
 
 * Adhere to [PEP 8](https://www.python.org/dev/peps/pep-0008/)
 * Use [pylint](https://www.pylint.org/)/[flake8](https://flake8.pycqa.org/) and [black](https://black.readthedocs.io/) to ensure as clean and well-formatted code as possible

README.md

@@ -2,12 +2,12 @@
 
 ![logo](./assets/identity/graphnet-logo-and-wordmark.png)
 
-| Usage | Development |
-| --- | --- |
-| [![status](https://joss.theoj.org/papers/eecab02fb1ecd174a5273750c1ea0baf/status.svg)](https://joss.theoj.org/papers/eecab02fb1ecd174a5273750c1ea0baf) | ![build](https://github.com/graphnet-team/graphnet/actions/workflows/build.yml/badge.svg) |
-| [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.6720188.svg)](https://doi.org/10.5281/zenodo.6720188) | ![code-quality](https://github.com/graphnet-team/graphnet/actions/workflows/code-quality.yml/badge.svg) |
-| [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) | [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) |
-| ![Supported python versions](https://img.shields.io/badge/python-3.8%20%7C%203.9%20%7C%203.10-blue) | [![Maintainability](https://api.codeclimate.com/v1/badges/b273a774112e32643162/maintainability)](https://codeclimate.com/github/graphnet-team/graphnet/maintainability) |
+| Usage                                                                                                                                                              | Development |
+|--------------------------------------------------------------------------------------------------------------------------------------------------------------------| --- |
+| [![status](https://joss.theoj.org/papers/eecab02fb1ecd174a5273750c1ea0baf/status.svg)](https://joss.theoj.org/papers/eecab02fb1ecd174a5273750c1ea0baf)             | ![build](https://github.com/graphnet-team/graphnet/actions/workflows/build.yml/badge.svg) |
+| [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.6720188.svg)](https://doi.org/10.5281/zenodo.6720188)                                                          | ![code-quality](https://github.com/graphnet-team/graphnet/actions/workflows/code-quality.yml/badge.svg) |
+| [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)                                               | [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) |
+| ![Supported python versions](https://img.shields.io/badge/python-3.8%20%7C%203.9%20%7C%203.10%20%7C%203.11-blue)                                                   | [![Maintainability](https://api.codeclimate.com/v1/badges/b273a774112e32643162/maintainability)](https://codeclimate.com/github/graphnet-team/graphnet/maintainability) |
 | [![Docker image](https://img.shields.io/docker/v/asogaard/graphnet?color=blue&logo=docker&sort=semver)](https://hub.docker.com/repository/docker/asogaard/graphnet) | [![Test Coverage](https://api.codeclimate.com/v1/badges/b273a774112e32643162/test_coverage)](https://codeclimate.com/github/graphnet-team/graphnet/test_coverage) |
 
 </center>
@@ -40,7 +40,7 @@ You may want `graphnet` to be able to interface with IceTray, e.g., when convert
 To achieve this, we recommend running the following commands in a clean bash shell:
 ```bash
 $ eval `/cvmfs/icecube.opensciencegrid.org/py3-v4.2.1/setup.sh`
-$ /cvmfs/icecube.opensciencegrid.org/py3-v4.2.1/Ubuntu_20.04_x86_64/metaprojects/icetray/v1.5.1/env-shell.sh
+$ /cvmfs/icecube.opensciencegrid.org/py3-v4.2.1/RHEL_7_x86_64/metaprojects/icetray/v1.5.1/env-shell.sh
 ```
 Optionally, you can alias these commands or save them as a bash script for convenience, as you will have to run these commands every time you want to use IceTray (with `graphnet`) in a clean shell.
 

configs/datasets/dev_lvl7_robustness_muon_neutrino_0000.yml

@@ -5,7 +5,7 @@ graph_definition:
     detector:
       arguments: {}
       class_name: IceCube86
-    dtype: null
+    dtype: torch.float32
     nb_nearest_neighbours: 8
     node_definition:
       arguments: {}

configs/datasets/test_data_sqlite.yml

@@ -5,7 +5,7 @@ graph_definition:
     detector:
       arguments: {}
       class_name: Prometheus
-    dtype: null
+    dtype: torch.float32
     nb_nearest_neighbours: 8
     node_definition:
       arguments: {}

configs/datasets/training_classification_example_data_sqlite.yml

@@ -5,7 +5,7 @@ graph_definition:
     detector:
       arguments: {}
       class_name: Prometheus
-    dtype: null
+    dtype: torch.float32
     nb_nearest_neighbours: 8
     node_definition:
       arguments: {}

configs/datasets/training_example_data_parquet.yml

@@ -5,7 +5,7 @@ graph_definition:
     detector:
       arguments: {}
       class_name: Prometheus
-    dtype: null
+    dtype: torch.float32
     nb_nearest_neighbours: 8
     node_definition:
       arguments: {}

configs/datasets/training_example_data_sqlite.yml

@@ -5,7 +5,7 @@ graph_definition:
     detector:
       arguments: {}
       class_name: Prometheus
-    dtype: null
+    dtype: torch.float32
     nb_nearest_neighbours: 8
     node_definition:
       arguments: {}

docker/gnn-benchmarking/apply.py

@@ -11,13 +11,17 @@
 
 from graphnet.deployment.i3modules import I3InferenceModule
 from graphnet.data.extractors.i3featureextractor import (
-    I3FeatureExtractorIceCubeDeepCore,
+    I3FeatureExtractorIceCubeUpgrade,
 )
 from graphnet.data.constants import FEATURES
+from graphnet.constants import PRETRAINED_MODEL_DIR
 
 
-# Constants (from Dockerfile)
-MODEL_PATH = "model.pth"
+# Constants
+MODEL_NAME = "total_neutrino_energy"
+BASE_PATH = f"{PRETRAINED_MODEL_DIR}/icecube/upgrade/QUESO"
+MODEL_CONFIG = f"{BASE_PATH}/{MODEL_NAME}/{MODEL_NAME}_config.yml"
+STATE_DICT = f"{BASE_PATH}/{MODEL_NAME}/{MODEL_NAME}_state_dict.pth"
 
 
 def main(
@@ -42,14 +46,15 @@ def main(
     input_files = [p for p in input_files if gcd_pattern not in p]
 
     # Construct I3InferenceModule(s)
-    extractor = I3FeatureExtractorIceCubeDeepCore(pulsemap=pulsemap)
+    extractor = I3FeatureExtractorIceCubeUpgrade(pulsemap=pulsemap)
 
     deployment_modules = [
         I3InferenceModule(
             pulsemap=pulsemap,
             features=FEATURES.DEEPCORE,
             pulsemap_extractor=extractor,
-            model=MODEL_PATH,
+            model_config=MODEL_CONFIG,
+            state_dict=STATE_DICT,
             gcd_file=gcd_file,
             prediction_columns=["energy_pred"],
             model_name="graphnet_dynedge_energy_reconstruction",

docker/gnn-benchmarking/dockerfile

@@ -17,6 +17,5 @@ WORKDIR /root
 
 ADD module.py ./module.py
 ADD apply.py ./apply.py
-ADD test_model.pth ./model.pth
 
 CMD [ "/bin/bash" ]

examples/01_icetray/03_i3_deployer_example.py

@@ -7,6 +7,7 @@
 from graphnet.constants import (
     TEST_DATA_DIR,
     EXAMPLE_OUTPUT_DIR,
+    PRETRAINED_MODEL_DIR,
 )
 from graphnet.data.constants import FEATURES, TRUTH
 from graphnet.data.extractors.i3featureextractor import (
@@ -34,7 +35,10 @@ def main() -> None:
     # configure input files, output folders and pulsemap
     pulsemap = "SplitInIcePulses"
     input_folders = [f"{TEST_DATA_DIR}/i3/upgrade_genie_step4_140028_000998"]
-    mock_model_path = f"{TEST_DATA_DIR}/models/mock_energy_model.pth"
+    base_path = f"{PRETRAINED_MODEL_DIR}/icecube/upgrade/QUESO"
+    model_name = "total_neutrino_energy"
+    model_config = f"{base_path}/{model_name}/{model_name}_config.yml"
+    state_dict = f"{base_path}/{model_name}/{model_name}_state_dict.pth"
     output_folder = f"{EXAMPLE_OUTPUT_DIR}/i3_deployment/upgrade_03_04"
     gcd_file = f"{TEST_DATA_DIR}/i3/upgrade_genie_step4_140028_000998/GeoCalibDetectorStatus_ICUpgrade.v58.mixed.V0.i3.bz2"
     input_files = []
@@ -46,7 +50,8 @@ def main() -> None:
         pulsemap=pulsemap,
         features=features,
         pulsemap_extractor=I3FeatureExtractorIceCubeUpgrade(pulsemap=pulsemap),
-        model=mock_model_path,
+        model_config=model_config,
+        state_dict=state_dict,
         gcd_file=gcd_file,
         prediction_columns=["energy"],
         model_name="graphnet_deployment_example",

examples/01_icetray/04_i3_module_in_native_icetray_example.py

@@ -11,6 +11,7 @@
 from graphnet.constants import (
     TEST_DATA_DIR,
     EXAMPLE_OUTPUT_DIR,
+    PRETRAINED_MODEL_DIR,
 )
 from graphnet.utilities.argparse import ArgumentParser
 from graphnet.utilities.imports import has_icecube_package
@@ -71,7 +72,10 @@ def main() -> None:
     # Constants
     features = FEATURES.UPGRADE
     input_folders = [f"{TEST_DATA_DIR}/i3/upgrade_genie_step4_140028_000998"]
-    mock_model_path = f"{TEST_DATA_DIR}/models/mock_energy_model.pth"
+    base_path = f"{PRETRAINED_MODEL_DIR}/icecube/upgrade/QUESO"
+    model_name = "total_neutrino_energy"
+    model_config = f"{base_path}/{model_name}/{model_name}_config.yml"
+    state_dict = f"{base_path}/{model_name}/{model_name}_state_dict.pth"
     output_folder = f"{EXAMPLE_OUTPUT_DIR}/i3_deployment/upgrade"
     gcd_file = f"{TEST_DATA_DIR}/i3/upgrade_genie_step4_140028_000998/GeoCalibDetectorStatus_ICUpgrade.v58.mixed.V0.i3.bz2"
     features = FEATURES.UPGRADE
@@ -84,7 +88,8 @@ def main() -> None:
         pulsemap=pulsemap,
         features=features,
         pulsemap_extractor=I3FeatureExtractorIceCubeUpgrade(pulsemap=pulsemap),
-        model=mock_model_path,
+        model_config=model_config,
+        state_dict=state_dict,
         gcd_file=gcd_file,
         prediction_columns=["energy"],
         model_name="graphnet_deployment_example",

examples/04_training/01_train_model.py

@@ -8,7 +8,7 @@
 from pytorch_lightning.utilities import rank_zero_only
 from graphnet.constants import EXAMPLE_OUTPUT_DIR
 from graphnet.data.dataloader import DataLoader
-from graphnet.models import Model
+from graphnet.models import StandardModel
 from graphnet.training.callbacks import ProgressBar
 from graphnet.utilities.argparse import ArgumentParser
 from graphnet.utilities.config import (
@@ -27,7 +27,6 @@ def main(
     early_stopping_patience: int,
     batch_size: int,
     num_workers: int,
-    prediction_names: Optional[List[str]],
     suffix: Optional[str] = None,
     wandb: bool = False,
 ) -> None:
@@ -49,7 +48,7 @@ def main(
 
     # Build model
     model_config = ModelConfig.load(model_config_path)
-    model = Model.from_config(model_config, trust=True)
+    model: StandardModel = StandardModel.from_config(model_config, trust=True)
 
     # Configuration
     config = TrainingConfig(
@@ -102,39 +101,31 @@ def main(
         **config.fit,
     )
 
+    # Save model to file
+    db_name = dataset_config.path.split("/")[-1].split(".")[0]
+    path = os.path.join(archive, db_name, run_name)
+    os.makedirs(path, exist_ok=True)
+    logger.info(f"Writing results to {path}")
+    model.save_state_dict(f"{path}/state_dict.pth")
+    model.save(f"{path}/model.pth")
+
     # Get predictions
     if isinstance(config.target, str):
-        prediction_columns = [config.target + "_pred"]
         additional_attributes = [config.target]
     else:
-        prediction_columns = [target + "_pred" for target in config.target]
         additional_attributes = config.target
 
-    if prediction_names:
-        prediction_columns = prediction_names
-
     logger.info(f"config.target: {config.target}")
-    logger.info(f"prediction_columns: {prediction_columns}")
+    logger.info(f"prediction_columns: {model.prediction_labels}")
 
     results = model.predict_as_dataframe(
         dataloaders["test"],
-        prediction_columns=prediction_columns,
         additional_attributes=additional_attributes + ["event_no"],
     )
-
-    # Save predictions and model to file
-    db_name = dataset_config.path.split("/")[-1].split(".")[0]
-    path = os.path.join(archive, db_name, run_name)
-    logger.info(f"Writing results to {path}")
-    os.makedirs(path, exist_ok=True)
-
     results.to_csv(f"{path}/results.csv")
-    model.save_state_dict(f"{path}/state_dict.pth")
-    model.save(f"{path}/model.pth")
 
 
 if __name__ == "__main__":
-
     # Parse command-line arguments
     parser = ArgumentParser(
         description="""
@@ -152,13 +143,6 @@ def main(
         "num-workers",
     )
 
-    parser.add_argument(
-        "--prediction-names",
-        nargs="+",
-        help="Names of each prediction output feature (default: %(default)s)",
-        default=None,
-    )
-
     parser.add_argument(
         "--suffix",
         type=str,
@@ -182,7 +166,6 @@ def main(
         args.early_stopping_patience,
         args.batch_size,
         args.num_workers,
-        args.prediction_names,
         args.suffix,
         args.wandb,
     )