Merge pull request graphnet-team#558 from RasmusOrsoe/graph_definitio…

…n_debug Graph Definition
RasmusOrsoe · Aug 14, 2023 · 2f444cf · 2f444cf
2 parents 3177d67 + 930e972
commit 2f444cf
Show file tree

Hide file tree

Showing 51 changed files with 1,331 additions and 801 deletions.
diff --git a/GETTING_STARTED.md b/GETTING_STARTED.md
diff --git a/configs/datasets/dev_lvl7_robustness_muon_neutrino_0000.yml b/configs/datasets/dev_lvl7_robustness_muon_neutrino_0000.yml
@@ -1,4 +1,17 @@
 path: /groups/icecube/asogaard/data/example/dev_lvl7_robustness_muon_neutrino_0000.db
+graph_definition:
+  arguments:
+    columns: [0, 1, 2]
+    detector:
+      arguments: {}
+      class_name: IceCube86
+    dtype: null
+    nb_nearest_neighbours: 8
+    node_definition:
+      arguments: {}
+      class_name: NodesAsPulses
+    node_feature_names: [dom_x, dom_y, dom_z, dom_time, charge, rde, pmt_area]
+  class_name: KNNGraph
 pulsemaps:
   - SRTTWOfflinePulsesDC
 features:

diff --git a/configs/datasets/test_data_sqlite.yml b/configs/datasets/test_data_sqlite.yml
@@ -1,26 +1,34 @@
-path: $GRAPHNET/data/tests/sqlite/oscNext_genie_level7_v02/oscNext_genie_level7_v02_first_5_frames.db
-pulsemaps:
-  - SRTInIcePulses
-features:
-  - dom_x
-  - dom_y
-  - dom_z
-  - dom_time
-  - charge
-  - rde
-  - pmt_area
-truth:
-  - energy
-  - position_x
-  - position_y
-  - position_z
-  - azimuth
-  - zenith
-  - pid
-  - elasticity
-  - sim_type
-  - interaction_type
+features: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+graph_definition:
+  arguments:
+    columns: [0, 1, 2]
+    detector:
+      arguments: {}
+      class_name: Prometheus
+    dtype: null
+    nb_nearest_neighbours: 8
+    node_definition:
+      arguments: {}
+      class_name: NodesAsPulses
+    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+  class_name: KNNGraph
 index_column: event_no
-truth_table: truth
-seed: 21
-selection: null
+loss_weight_column: null
+loss_weight_default_value: null
+loss_weight_table: null
+node_truth: null
+node_truth_table: null
+path: $GRAPHNET/data/examples/sqlite/prometheus/prometheus-events.db
+pulsemaps: total
+seed: null
+selection: null
+string_selection: null
+truth: [injection_energy, injection_type, injection_interaction_type, injection_zenith,
+  injection_azimuth, injection_bjorkenx, injection_bjorkeny, injection_position_x,
+  injection_position_y, injection_position_z, injection_column_depth, primary_lepton_1_type,
+  primary_hadron_1_type, primary_lepton_1_position_x, primary_lepton_1_position_y,
+  primary_lepton_1_position_z, primary_hadron_1_position_x, primary_hadron_1_position_y,
+  primary_hadron_1_position_z, primary_lepton_1_direction_theta, primary_lepton_1_direction_phi,
+  primary_hadron_1_direction_theta, primary_hadron_1_direction_phi, primary_lepton_1_energy,
+  primary_hadron_1_energy, total_energy]
+truth_table: mc_truth
diff --git a/configs/datasets/training_classification_example_data_sqlite.yml b/configs/datasets/training_classification_example_data_sqlite.yml
@@ -1,4 +1,17 @@
 path: $GRAPHNET/data/examples/sqlite/prometheus/prometheus-events.db
+graph_definition:
+  arguments:
+    columns: [0, 1, 2]
+    detector:
+      arguments: {}
+      class_name: Prometheus
+    dtype: null
+    nb_nearest_neighbours: 8
+    node_definition:
+      arguments: {}
+      class_name: NodesAsPulses
+    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+  class_name: KNNGraph
 pulsemaps:
   - total
 features:

diff --git a/configs/datasets/training_example_data_parquet.yml b/configs/datasets/training_example_data_parquet.yml
@@ -1,4 +1,17 @@
 path: $GRAPHNET/data/examples/parquet/prometheus/prometheus-events.parquet
+graph_definition:
+  arguments:
+    columns: [0, 1, 2]
+    detector:
+      arguments: {}
+      class_name: Prometheus
+    dtype: null
+    nb_nearest_neighbours: 8
+    node_definition:
+      arguments: {}
+      class_name: NodesAsPulses
+    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+  class_name: KNNGraph
 pulsemaps:
   - total
 features:

diff --git a/configs/datasets/training_example_data_sqlite.yml b/configs/datasets/training_example_data_sqlite.yml
@@ -1,4 +1,17 @@
 path: $GRAPHNET/data/examples/sqlite/prometheus/prometheus-events.db
+graph_definition:
+  arguments:
+    columns: [0, 1, 2]
+    detector:
+      arguments: {}
+      class_name: Prometheus
+    dtype: null
+    nb_nearest_neighbours: 8
+    node_definition:
+      arguments: {}
+      class_name: NodesAsPulses
+    node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+  class_name: KNNGraph
 pulsemaps:
   - total
 features:

diff --git a/configs/models/dynedge_PID_classification_example.yml b/configs/models/dynedge_PID_classification_example.yml
@@ -1,14 +1,4 @@
 arguments:
-  coarsening: null
-  detector:
-    ModelConfig:
-      arguments:
-        graph_builder:
-          ModelConfig:
-            arguments: {columns: null, nb_nearest_neighbours: 8}
-            class_name: KNNGraphBuilder
-        scalers: null
-      class_name: Prometheus
   gnn:
     ModelConfig:
       arguments:
@@ -21,11 +11,29 @@ arguments:
         post_processing_layer_sizes: null
         readout_layer_sizes: null
       class_name: DynEdge
+  graph_definition:
+    ModelConfig:
+      arguments:
+        columns: [0, 1, 2]
+        detector:
+          ModelConfig:
+            arguments: {}
+            class_name: Prometheus
+        dtype: null
+        nb_nearest_neighbours: 8
+        node_definition:
+          ModelConfig:
+            arguments: {}
+            class_name: NodesAsPulses
+        node_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: 1e-03, lr: 1e-05}
-  scheduler_class: '!class torch.optim.lr_scheduler ReduceLROnPlateau'
-  scheduler_config: {frequency: 1, monitor: val_loss}
-  scheduler_kwargs: {patience: 1}
+  optimizer_kwargs: {eps: 0.001, lr: 0.001}
+  scheduler_class: '!class graphnet.training.callbacks PiecewiseLinearLR'
+  scheduler_config: {interval: step}
+  scheduler_kwargs:
+    factors: [0.01, 1, 0.01]
+    milestones: [0, 20.0, 80]
   tasks:
   - ModelConfig:
       arguments:

diff --git a/configs/models/dynedge_energy_example.yml b/configs/models/dynedge_energy_example.yml
diff --git a/configs/models/dynedge_position_custom_scaling_example.yml b/configs/models/dynedge_position_custom_scaling_example.yml
@@ -1,14 +1,24 @@
 arguments:
-  coarsening: null
-  detector:
+  graph_definition:
     ModelConfig:
       arguments:
-        graph_builder:
+        detector:
           ModelConfig:
-            arguments: {columns: null, nb_nearest_neighbours: 8}
-            class_name: KNNGraphBuilder
-        scalers: null
-      class_name: IceCubeDeepCore
+            arguments: {}
+            class_name: Prometheus
+        dtype: null
+        edge_definition:
+          ModelConfig:
+            arguments:
+              columns: [0, 1, 2]
+              nb_nearest_neighbours: 8
+            class_name: KNNEdges
+        node_definition:
+          ModelConfig:
+            arguments: {}
+            class_name: NodesAsPulses
+        node_feature_names: null
+      class_name: KNNGraph
   gnn:
     ModelConfig:
       arguments:

diff --git a/configs/models/example_direction_reconstruction_model.yml b/configs/models/example_direction_reconstruction_model.yml
@@ -1,14 +1,20 @@
 arguments:
-  coarsening: null
-  detector:
+  graph_definition:
     ModelConfig:
       arguments:
-        graph_builder:
+        columns: [0, 1, 2]
+        detector:
           ModelConfig:
-            arguments: {columns: null, nb_nearest_neighbours: 8}
-            class_name: KNNGraphBuilder
-        scalers: null
-      class_name: Prometheus
+            arguments: {}
+            class_name: Prometheus
+        dtype: null
+        nb_nearest_neighbours: 8
+        node_definition:
+          ModelConfig:
+            arguments: {}
+            class_name: NodesAsPulses
+        node_feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t]
+      class_name: KNNGraph
   gnn:
     ModelConfig:
       arguments:

diff --git a/configs/models/example_energy_reconstruction_model.yml b/configs/models/example_energy_reconstruction_model.yml
@@ -1,14 +1,4 @@
 arguments:
-  coarsening: null
-  detector:
-    ModelConfig:
-      arguments:
-        graph_builder:
-          ModelConfig:
-            arguments: {columns: null, nb_nearest_neighbours: 8}
-            class_name: KNNGraphBuilder
-        scalers: null
-      class_name: Prometheus
   gnn:
     ModelConfig:
       arguments:
@@ -21,11 +11,29 @@ arguments:
         post_processing_layer_sizes: null
         readout_layer_sizes: null
       class_name: DynEdge
+  graph_definition:
+    ModelConfig:
+      arguments:
+        columns: [0, 1, 2]
+        detector:
+          ModelConfig:
+            arguments: {}
+            class_name: Prometheus
+        dtype: null
+        nb_nearest_neighbours: 8
+        node_definition:
+          ModelConfig:
+            arguments: {}
+            class_name: NodesAsPulses
+        node_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: 1e-05}
-  scheduler_class: '!class torch.optim.lr_scheduler ReduceLROnPlateau'
-  scheduler_config: {frequency: 1, monitor: val_loss}
-  scheduler_kwargs: {patience: 5}
+  optimizer_kwargs: {eps: 0.001, lr: 0.001}
+  scheduler_class: '!class graphnet.training.callbacks PiecewiseLinearLR'
+  scheduler_config: {interval: step}
+  scheduler_kwargs:
+    factors: [0.01, 1, 0.01]
+    milestones: [0, 20.0, 80]
   tasks:
   - ModelConfig:
       arguments:
@@ -35,8 +43,9 @@ arguments:
             arguments: {}
             class_name: LogCoshLoss
         loss_weight: null
+        prediction_labels: null
         target_labels: total_energy
-        transform_inference: null
+        transform_inference: '!lambda x: torch.pow(10,x)'
         transform_prediction_and_target: '!lambda x: torch.log10(x)'
         transform_support: null
         transform_target: null

diff --git a/configs/models/example_vertex_position_reconstruction_model.yml b/configs/models/example_vertex_position_reconstruction_model.yml
@@ -1,14 +1,4 @@
 arguments:
-  coarsening: null
-  detector:
-    ModelConfig:
-      arguments:
-        graph_builder:
-          ModelConfig:
-            arguments: {columns: null, nb_nearest_neighbours: 8}
-            class_name: KNNGraphBuilder
-        scalers: null
-      class_name: Prometheus
   gnn:
     ModelConfig:
       arguments:
@@ -21,6 +11,22 @@ arguments:
         post_processing_layer_sizes: null
         readout_layer_sizes: null
       class_name: DynEdge
+  graph_definition:
+    ModelConfig:
+      arguments:
+        columns: [0, 1, 2]
+        detector:
+          ModelConfig:
+            arguments: {}
+            class_name: Prometheus
+        dtype: null
+        nb_nearest_neighbours: 8
+        node_definition:
+          ModelConfig:
+            arguments: {}
+            class_name: NodesAsPulses
+        node_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}
   scheduler_class: '!class graphnet.training.callbacks PiecewiseLinearLR'

diff --git a/examples/01_icetray/02_compare_sqlite_and_parquet.py b/examples/01_icetray/02_compare_sqlite_and_parquet.py
@@ -4,11 +4,9 @@
 
 from graphnet.constants import EXAMPLE_OUTPUT_DIR, TEST_DATA_DIR
 from graphnet.data.constants import FEATURES, TRUTH
-from graphnet.data.sqlite import (
-    SQLiteDataConverter,
-    SQLiteDataset,
-)
-from graphnet.data.parquet import ParquetDataConverter, ParquetDataset
+from graphnet.data.sqlite import SQLiteDataConverter
+from graphnet.data.parquet import ParquetDataConverter
+from graphnet.data.dataset import SQLiteDataset, ParquetDataset
 from graphnet.data.extractors import (
     I3FeatureExtractorIceCube86,
     I3TruthExtractor,