Merge branch 'paper-test-branch' into new_dataset

src/graphnet/data/curated_datamodule.py

@@ -285,7 +285,17 @@ def prepare_data(self) -> None:
             return
         else:
             # Download, unzip and delete zipped file
+
             os.makedirs(self.dataset_dir, exist_ok=True)
+
+            os.makedirs(self.dataset_dir)
+            _, file_name = os.path.split(file_hash)
+            extension = ".tar.gz" if ".tar.gz" not in file_name else ""
+            file_path = os.path.join(
+                self.dataset_dir,
+                file_name + extension,
+            )
+
             os.system(f"wget -O {file_path} {self._mirror}/{file_hash}")
             print("Unzipping file, this might take a while..")
             if self._backend == "parquet":

src/graphnet/models/detector/detector.py

@@ -33,7 +33,7 @@ def forward(  # type: ignore
     @property
     def geometry_table(self) -> pd.DataFrame:
         """Public get method for retrieving a `Detector`s geometry table."""
-        if ~hasattr(self, "_geometry_table"):
+        if not hasattr(self, "_geometry_table"):
             try:
                 assert hasattr(self, "geometry_table_path")
             except AssertionError as e:
@@ -60,6 +60,16 @@ def sensor_index_name(self) -> str:
         """Public get method for retrieving the sensor id column name."""
         return self.sensor_id_column
 
+    @property
+    def sensor_time_name(self) -> str:
+        """Public get method for retrieving the sensor time column name."""
+        return self.sensor_time_column
+
+    @property
+    def charge_name(self) -> str:
+        """Public get method for retrieving the charge column name."""
+        return self.charge_column
+
     @final
     def _standardize(
         self, input_features: torch.tensor, input_feature_names: List[str]

src/graphnet/models/detector/icecube.py

@@ -17,6 +17,8 @@ class IceCube86(Detector):
     xyz = ["dom_x", "dom_y", "dom_z"]
     string_id_column = "string"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "dom_time"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension of input data."""
@@ -57,6 +59,8 @@ class IceCubeKaggle(Detector):
     xyz = ["x", "y", "z"]
     string_id_column = "string"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "time"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension of input data."""
@@ -122,6 +126,8 @@ class IceCubeUpgrade(Detector):
     xyz = ["dom_x", "dom_y", "dom_z"]
     string_id_column = "string"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "dom_time"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension of input data."""

src/graphnet/models/detector/liquido.py

@@ -17,6 +17,8 @@ class LiquidO_v1(Detector):
     xyz = ["sipm_x", "sipm_y", "sipm_z"]
     string_id_column = "fiber_id"
     sensor_id_column = "sipm_id"
+    sensor_time_column = "t"
+    charge_column = None
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""

src/graphnet/models/detector/prometheus.py

@@ -17,6 +17,8 @@ class ORCA150SuperDense(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -47,6 +49,8 @@ class TRIDENT1211(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -77,6 +81,8 @@ class IceCubeUpgrade7(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -107,6 +113,8 @@ class WaterDemo81(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -137,6 +145,8 @@ class BaikalGVD8(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -167,6 +177,8 @@ class IceDemo81(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -197,6 +209,8 @@ class ARCA115(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -227,6 +241,8 @@ class ORCA150(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -257,6 +273,8 @@ class IceCube86Prometheus(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -287,6 +305,8 @@ class IceCubeDeepCore8(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -317,6 +337,8 @@ class IceCubeGen2(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""
@@ -344,6 +366,8 @@ class PONETriangle(Detector):
     xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"]
     string_id_column = "sensor_string_id"
     sensor_id_column = "sensor_id"
+    sensor_time_column = "t"
+    charge_column = "charge"
 
     def feature_map(self) -> Dict[str, Callable]:
         """Map standardization functions to each dimension."""