Merge branch 'graphnet-team:main' into fix-tito

.github/workflows/build.yml

@@ -36,8 +36,16 @@ jobs:
     name: Unit tests - IceTray
     needs: [ check-codeclimate-credentials ]
     runs-on: ubuntu-latest
-    container: icecube/icetray:icetray-prod-v1.8.1-ubuntu20.04-X64
+    container:
+      image: icecube/icetray:icetray-prod-v1.8.1-ubuntu20.04-X64
+      options: --user root
     steps:
+      - name: install git
+        run: |
+          apt-get --yes install sudo
+          sudo apt update --fix-missing --yes
+          sudo apt upgrade --yes
+          sudo apt-get install --yes git-all
       - name: Set environment variables
         run: |
           echo "PATH=/usr/local/icetray/bin:$PATH" >> $GITHUB_ENV
@@ -57,13 +65,13 @@ jobs:
           editable: true
       - name: Run unit tests and generate coverage report
         run: |
-          coverage run --source=graphnet -m pytest tests/ --ignore=tests/examples
-          coverage run --source=graphnet -m pytest tests/examples/01_icetray
+          coverage run --source=graphnet -m pytest tests/ --ignore=tests/examples/04_training
+          coverage run -a --source=graphnet -m pytest tests/examples/04_training
           coverage xml -o coverage.xml
     
-      #- name: Work around permission issue
-      #  run: |
-      #    git config --global --add safe.directory /__w/graphnet/graphnet
+      - name: Work around permission issue
+        run: |
+          git config --global --add safe.directory /__w/graphnet/graphnet
       - name: Publish code coverage
         uses: paambaati/[email protected]
         if: needs.check-codeclimate-credentials.outputs.has_credentials == 'true'
@@ -93,14 +101,21 @@ jobs:
           editable: true
       - name: Print available disk space after graphnet install
         run: df -h
+      - name: Print packages in pip
+        run: |
+          pip show torch
+          pip show torch-geometric
+          pip show torch-cluster
+          pip show torch-sparse
+          pip show torch-scatter
       - 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/utilities --ignore=tests/data/ --ignore=tests/deployment/ --ignore=tests/examples/01_icetray/
-          coverage run --source=graphnet -m pytest tests/utilities
+          coverage run -a --source=graphnet -m pytest tests/utilities
           coverage report -m
       - name: Print available disk space after unit tests
-        run: df -h  
+        run: df -h
 
   build-macos:
     name: Unit tests - macOS
@@ -116,8 +131,15 @@ jobs:
         with:
           editable: true
           hardware: "macos"
+      - name: Print packages in pip
+        run: |
+            pip show torch
+            pip show torch-geometric
+            pip show torch-cluster
+            pip show torch-sparse
+            pip show torch-scatter
       - 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/ --ignore=tests/deployment/ --ignore=tests/examples/
-          coverage report -m
+          coverage report -m

.github/workflows/docker.yml

@@ -19,6 +19,28 @@ jobs:
       packages: write
       contents: read
     steps:
+      - name: Before Clean-up
+        run: |
+          echo "Free space:"
+          df -h
+
+      - name: Free Disk Space
+        uses: jlumbroso/free-disk-space@main
+        with:
+          tool-cache: true
+
+          # all of these default to true, but feel free to set to
+          # false if necessary for your workflow
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: true
+          swap-storage: true
+
+      - name: After Clean-up
+        run: |
+          echo "Free space:"
+          df -h
       - name: Checkout
         uses: actions/checkout@v3
       - name: Set up QEMU

README.md

@@ -82,7 +82,6 @@ $ conda create --name graphnet python=3.8 gcc_linux-64 gxx_linux-64 libgcc cudat
 $ conda activate graphnet  # Optional
 (graphnet) $ pip install -r requirements/torch_cpu.txt -e .[develop,torch]  # CPU-only torch
 (graphnet) $ pip install -r requirements/torch_gpu.txt -e .[develop,torch]  # GPU support
-(graphnet) $ pip install -r requirements/torch_macos.txt -e .[develop,torch]  # On macOS
 ```
 This should allow you to e.g. run the scripts in [examples/](./examples/) out of the box.
 

configs/datasets/dev_lvl7_robustness_muon_neutrino_0000.yml

@@ -10,7 +10,7 @@ graph_definition:
     node_definition:
       arguments: {}
       class_name: NodesAsPulses
-    node_feature_names: [dom_x, dom_y, dom_z, dom_time, charge, rde, pmt_area]
+    input_feature_names: [dom_x, dom_y, dom_z, dom_time, charge, rde, pmt_area]
   class_name: KNNGraph
 pulsemaps:
   - SRTTWOfflinePulsesDC

configs/datasets/test_data_sqlite.yml

@@ -10,7 +10,7 @@ graph_definition:
     node_definition:
       arguments: {}
       class_name: NodesAsPulses
-    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+    input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
   class_name: KNNGraph
 index_column: event_no
 loss_weight_column: null

configs/datasets/training_classification_example_data_sqlite.yml

@@ -10,7 +10,7 @@ graph_definition:
     node_definition:
       arguments: {}
       class_name: NodesAsPulses
-    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+    input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
   class_name: KNNGraph
 pulsemaps:
   - total

configs/datasets/training_example_data_parquet.yml

@@ -10,7 +10,7 @@ graph_definition:
     node_definition:
       arguments: {}
       class_name: NodesAsPulses
-    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+    input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
   class_name: KNNGraph
 pulsemaps:
   - total

configs/datasets/training_example_data_sqlite.yml

@@ -10,7 +10,7 @@ graph_definition:
     node_definition:
       arguments: {}
       class_name: NodesAsPulses
-    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+    input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
   class_name: KNNGraph
 pulsemaps:
   - total

configs/models/dynedge_PID_classification_example.yml

@@ -25,7 +25,7 @@ arguments:
           ModelConfig:
             arguments: {}
             class_name: NodesAsPulses
-        node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+        input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
       class_name: KNNGraph
   optimizer_class: '!class torch.optim.adam Adam'
   optimizer_kwargs: {eps: 0.001, lr: 0.001}

configs/models/dynedge_position_custom_scaling_example.yml

@@ -17,7 +17,7 @@ arguments:
           ModelConfig:
             arguments: {}
             class_name: NodesAsPulses
-        node_feature_names: null
+        input_feature_names: null
       class_name: KNNGraph
   gnn:
     ModelConfig:

configs/models/example_direction_reconstruction_model.yml

@@ -13,7 +13,7 @@ arguments:
           ModelConfig:
             arguments: {}
             class_name: NodesAsPulses
-        node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+        input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
       class_name: KNNGraph
   gnn:
     ModelConfig:

configs/models/example_energy_reconstruction_model.yml

@@ -25,7 +25,7 @@ arguments:
           ModelConfig:
             arguments: {}
             class_name: NodesAsPulses
-        node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+        input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
       class_name: KNNGraph
   optimizer_class: '!class torch.optim.adam Adam'
   optimizer_kwargs: {eps: 0.001, lr: 0.001}

configs/models/example_vertex_position_reconstruction_model.yml

@@ -25,7 +25,7 @@ arguments:
           ModelConfig:
             arguments: {}
             class_name: NodesAsPulses
-        node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+        input_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
       class_name: KNNGraph
   optimizer_class: '!class torch.optim.adam Adam'
   optimizer_kwargs: {eps: 0.001, lr: 0.001}

examples/02_data/04_ensemble_dataset.py

@@ -24,7 +24,7 @@
     detector=IceCubeDeepCore(),
     node_definition=NodesAsPulses(),
     nb_nearest_neighbours=8,
-    node_feature_names=features,
+    input_feature_names=features,
 )
 
 

examples/05_pisa/02_make_pipeline_database.py

@@ -65,7 +65,7 @@ def main() -> None:
         detector=IceCubeDeepCore(),
         node_definition=NodesAsPulses(),
         nb_nearest_neighbours=8,
-        node_feature_names=FEATURES.DEEPCORE,
+        input_feature_names=FEATURES.DEEPCORE,
     )
 
     # Remove `interaction_time` if it exists

requirements/torch_cpu.txt

@@ -1,2 +1,2 @@
 --find-links https://download.pytorch.org/whl/cpu
---find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html
+--find-links https://data.pyg.org/whl/torch-2.1.0+cpu.html

requirements/torch_gpu.txt

@@ -1,4 +1,4 @@
-# Contains packages recommended for functional performance
+# Contains packages requirements for GPU installation
 --find-links https://download.pytorch.org/whl/torch_stable.html
-torch==2.0.1+cu117
---find-links https://data.pyg.org/whl/torch-2.0.0+cu117.html
+torch==2.1.0+cu118
+--find-links https://data.pyg.org/whl/torch-2.1.0+cu118.html

requirements/torch_macos.txt

@@ -1,2 +1,2 @@
 --find-links https://download.pytorch.org/whl/torch_stable.html
---find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html
+--find-links https://data.pyg.org/whl/torch-2.1.0+cpu.html

setup.py

@@ -47,11 +47,11 @@
         "versioneer",
     ],
     "torch": [
-        "torch>=2.0",
+        "torch>=2.1",
         "torch-cluster>=1.6",
         "torch-scatter>=2.0",
         "torch-sparse>=0.6",
-        "torch-geometric>=2.1",
+        "torch-geometric>=2.3",
         "pytorch-lightning>=2.0",
     ],
 }

src/graphnet/data/__init__.py

@@ -3,3 +3,4 @@
 `graphnet.data` enables converting domain-specific data to industry-standard,
 intermediate file formats  and reading this data.
 """
+from .filters import I3Filter, I3FilterMask

src/graphnet/data/dataconverter.py

@@ -39,6 +39,7 @@
 from graphnet.utilities.filesys import find_i3_files
 from graphnet.utilities.imports import has_icecube_package
 from graphnet.utilities.logging import Logger
+from graphnet.data.filters import I3Filter, NullSplitI3Filter
 
 if has_icecube_package():
     from icecube import icetray, dataio  # pyright: reportMissingImports=false
@@ -107,6 +108,7 @@ def __init__(
         workers: int = 1,
         index_column: str = "event_no",
         icetray_verbose: int = 0,
+        i3_filters: List[I3Filter] = [],
     ):
         """Construct DataConverter.
 
@@ -167,6 +169,14 @@ def __init__(
         self._input_file_batch_pattern = input_file_batch_pattern
         self._workers = workers
 
+        # I3Filters (NullSplitI3Filter is always included)
+        self._i3filters = [NullSplitI3Filter()] + i3_filters
+
+        for filter in self._i3filters:
+            assert isinstance(
+                filter, I3Filter
+            ), f"{type(filter)} is not a subclass of I3Filter"
+
         # Create I3Extractors
         self._extractors = I3ExtractorCollection(*extractors)
 
@@ -433,6 +443,7 @@ def _extract_data(self, fileset: FileSet) -> List[OrderedDict]:
             except Exception as e:
                 if "I3" in str(e):
                     continue
+            # check if frame should be skipped
             if self._skip_frame(frame):
                 continue
 
@@ -555,14 +566,15 @@ def _get_output_file(self, input_file: str) -> str:
         return output_file
 
     def _skip_frame(self, frame: "icetray.I3Frame") -> bool:
-        """Check if frame should be skipped.
-
-        Args:
-            frame: I3Frame to check.
+        """Check the user defined filters.
 
         Returns:
-            True if frame is a null split frame, else False.
+            bool: True if frame should be skipped, False otherwise.
         """
-        if frame["I3EventHeader"].sub_event_stream == "NullSplit":
-            return True
-        return False
+        if self._i3filters is None:
+            return False  # No filters defined, so we keep the frame
+
+        for filter in self._i3filters:
+            if not filter(frame):
+                return True  # keep_frame call false, skip the frame.
+        return False  # All filter keep_frame calls true, keep the frame.

src/graphnet/data/dataset/dataset.py

@@ -282,7 +282,7 @@ def __init__(
         self._index_column = index_column
         self._truth_table = truth_table
         self._loss_weight_default_value = loss_weight_default_value
-        self._graph_definition = graph_definition
+        self._graph_definition = deepcopy(graph_definition)
 
         if node_truth is not None:
             assert isinstance(node_truth_table, str)
@@ -629,8 +629,8 @@ def _create_graph(
         # Construct graph data object
         assert self._graph_definition is not None
         graph = self._graph_definition(
-            node_features=node_features,
-            node_feature_names=self._features[
+            input_features=node_features,
+            input_feature_names=self._features[
                 1:
             ],  # first entry is index column
             truth_dicts=truth_dicts,